Compare commits

...
Author SHA1 Message Date
Robert Wojciechowski 2df73e4b9f feat: add image-list return format and quality controls 2026-10-14 00:44:45 +00:00
Robert Wojciechowski b14c51c365 Merge branch 'feat/image-list-serverless-base' into feat/image-list-serverless 2026-10-14 00:37:48 +00:00
Robert Wojciechowski 3699477f75 Merge branch 'feat/serverless-multimodal' into feat/image-list-serverless-base 2026-10-14 00:37:48 +00:00
Robert Wojciechowski b636fd5202 feat: add explicit image return codecs and quality controls 2026-10-14 00:37:42 +00:00
Robert Wojciechowski e9c95bf110 feat: schedule image-list items on bounded serverless slots 2026-10-14 00:16:54 +00:00
Robert Wojciechowski 1fb4d4a284 fix: preserve mixed-mode preflight across execution planner integration 2026-10-13 23:56:35 +00:00
Robert Wojciechowski 4e823397b7 feat: preserve required output utilities in image-list execution 2026-10-13 23:55:26 +00:00
Robert Wojciechowski f0927f1166 feat: add versioned lossless serverless media transport 2026-10-13 23:53:58 +00:00
Robert Wojciechowski 41caf02481 docs: tighten image-list and serverless wording 2026-10-13 22:38:10 +00:00
Robert Wojciechowski 494d5506a0 README: add Distributed Image List Start/End to nodes table 2026-10-13 22:27:07 +00:00
Robert Wojciechowski 7ed60b93ca feat: dynamically distribute image lists across direct workers (#98)
* feat: add Runpod serverless image collector pools

* fix: prevent worker settings clipping and remove preview copy

* fix: initialize worker settings before mounting the form

* fix: validate complete serverless requests against 10 MB limit

* fix: simplify serverless cards and keep status dots circular

* fix: match serverless endpoint box to master card

* fix: preserve serverless capacity and recovery across job failures

Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.

* feat: add dynamic per-image list queue with Start and End nodes

* fix(image-list): reject upstream data bypasses around Start

* fix(image-list): make optional End text a connection-only input

* fix(image-list): isolate unavailable workers during preflight

* fix(image-list): bound acknowledgements and configure run deadline

* fix(media-sync): preserve output and temp source roots

* fix(image-list): own job lifecycle and preserve worker preparation

* feat(image-list): let master participate unless orchestrator-only

* feat(image-list)!: replace fixed Start extras with growable value_N slots

Distributed Image List Start no longer has reference_image, text,
strength, seed or per-slot mode dropdowns. It now takes images, mask and
up to 10 generic value_N inputs (IMAGE, MASK, STRING, INT, FLOAT,
BOOLEAN) that grow via native V3 Autogrow, each with a matching output.

- One value is shared by every image; one per image pairs in order; any
  other count fails before dispatch. Only images is split into frames.
- Outputs are declared up front (image, mask, item_index, context,
  value_1..value_10). web/imageListStart.js hides unused value outputs,
  mirrors input types, and follows Autogrow compaction so downstream
  links keep their value.
- Image-list protocol bumped to v2; older workers fail preflight.

Breaking: saved workflows using the old Start inputs must recreate Start.

* fix(serverless): disable Clear saved API key while the worker is enabled

The checkbox read "Clear saved API key (disable pool first)", but the
panel never defines "pool" and saving with it ticked on an enabled worker
was rejected by the backend.

- Disable the checkbox while the worker is enabled, with a tooltip saying
  to turn the worker off first. Also disable it when no key is saved.
- Use "worker" instead of "pool" in user-facing serverless errors, the
  serverless guide and the README entry.

* feat(serverless): replace Concurrent jobs with a single Images per run setting

- Remove the separate Concurrent jobs field; rename Jobs per run to
  Images per run (1-16) with a tooltip explaining that Runpod's max
  workers decides how many GPUs run them.
- Backend uses a fixed per-worker limit of 16 outstanding jobs, so all
  images of a run are submitted at once. Unresolved jobs still count
  against it. The limit is constant, so it never has to change while
  jobs are unresolved.
- Legacy 'concurrency' in saved configs and older browser tabs is
  accepted and ignored/dropped; it no longer triggers the busy guard.
- Revert the enable toggle and show the reason when the backend rejects it.
- Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer
  endpoint-ID error otherwise.
- Confirm before a type change deletes a saved Runpod API key.
- Update docs and tests.
2026-10-13 15:20:05 -07:00
Robert Wojciechowski 6b94242110 feat: add Runpod serverless image Collector pools (#96)
* feat: add Runpod serverless image collector pools

* fix: prevent worker settings clipping and remove preview copy

* fix: initialize worker settings before mounting the form

* fix: validate complete serverless requests against 10 MB limit

* fix: simplify serverless cards and keep status dots circular

* fix: match serverless endpoint box to master card

* fix: preserve serverless capacity and recovery across job failures

Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.

* fix(serverless): disable Clear saved API key while the worker is enabled

The checkbox read "Clear saved API key (disable pool first)", but the
panel never defines "pool" and saving with it ticked on an enabled worker
was rejected by the backend.

- Disable the checkbox while the worker is enabled, with a tooltip saying
  to turn the worker off first. Also disable it when no key is saved.
- Use "worker" instead of "pool" in user-facing serverless errors, the
  serverless guide and the README entry.

* feat(serverless): replace Concurrent jobs with a single Images per run setting

- Remove the separate Concurrent jobs field; rename Jobs per run to
  Images per run (1-16) with a tooltip explaining that Runpod's max
  workers decides how many GPUs run them.
- Backend uses a fixed per-worker limit of 16 outstanding jobs, so all
  images of a run are submitted at once. Unresolved jobs still count
  against it. The limit is constant, so it never has to change while
  jobs are unresolved.
- Legacy 'concurrency' in saved configs and older browser tabs is
  accepted and ignored/dropped; it no longer triggers the busy guard.
- Revert the enable toggle and show the reason when the backend rejects it.
- Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer
  endpoint-ID error otherwise.
- Confirm before a type change deletes a saved Runpod API key.
- Update docs and tests.
2026-10-13 15:19:00 -07:00
40678a51c4 feat(worker): terminate pruned worker prompts with a hidden no-op sink (#99)
Worker prompts drop everything downstream of DistributedCollector /
UltimateSDUpscaleDistributed, so the master re-added a PreviewImage
(or PreviewAudio) to give ComfyUI an output node. That preview wrote
every result frame as a PNG into the worker's temp directory: for a
362-frame 1344x768 video, ~2 GB per run (RAM when temp is on tmpfs)
plus CPU time spent PNG-encoding frames nobody looks at.

- Add DistributedWorkerSink: an output node that discards its input.
  It is is_dev_only, so it stays out of node search and menus.
- The master probes each worker with GET /object_info/DistributedWorkerSink
  while preparing its prompt and uses the sink when it is registered.
  Workers on older versions (or a failed probe) keep the PreviewImage /
  PreviewAudio terminator, so mixed versions keep working.

Verified with a real master plus one updated and one old worker:
the updated worker ran EmptyImage -> Collector -> Sink and left its
temp directory empty; the old worker still received PreviewImage.


Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6

Co-authored-by: ymiyad <cab-13-might@icloud.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-14 07:27:19 +11:00
Robert Wojciechowski 9345a8e71b docs: link community add-ons and forks 2026-10-12 01:20:46 +00:00
Robert Wojciechowski a630f55bd9 docs: move compressed result transfer guide out of README 2026-10-11 20:55:37 +00:00
6edd3fda38 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>
2026-10-11 20:50:23 +00:00
Robert Wojciechowski 6f85129e45 Bump version from 1.4.7 to 1.5.0 2026-10-11 10:06:41 +11:00
Robert Wojciechowski a53d88f7fb test: remove redundant coverage and consolidate happy paths (#95) 2026-10-10 16:00:41 -07:00
Robert Wojciechowski 9be342f5fb feat: migrate distributed nodes to native ComfyUI V3 API (#94)
* feat: migrate distributed nodes to native ComfyUI V3 API

* fix: hide Distributed Value backing fields in Nodes 2.0

Use native widget visibility without converting backing fields into sockets. Preserve widget types, serializers and options identity; cover edits, reload, worker refresh and linked widgets.
2026-10-10 16:00:36 -07:00
Robert Wojciechowski a1faf7e916 fix: isolate local worker databases and avoid port conflicts (#93)
* fix: isolate local worker databases and avoid port conflicts

* fix: harden local worker startup preflight and database paths

Handle auto-launch before the master binds, allow TIME_WAIT restarts without accepting active listeners, resolve relative paths in the worker cwd, and round-trip SQLite URLs through SQLAlchemy. Add regressions and remove the README additions.
2026-10-10 16:00:31 -07:00
Robert Wojciechowski 32ac027e0e docs: align selected copy with current behavior 2026-07-12 11:13:22 +00:00
Robert Wojciechowski fec959e1ac chore: bump version to 1.4.7 2026-07-12 10:40:27 +00:00
Robert Wojciechowski cedc8d4591 feat: support audio-only distributed workflows (#87)
Support image, audio, or combined distributed collection with guarded audio-only worker, master, endpoint, preview, and delegate-only paths.
2026-07-12 20:36:43 +10:00
Robert Wojciechowski db01824d6e Update pyproject.toml 2026-06-06 09:10:43 +10:00
Robert Wojciechowski 6e7818bee6 fix: preserve delegate-only master config inputs (#86)
Fixes #85
2026-06-05 17:12:25 +10:00
Robert Wojciechowski c246d965e7 Update pyproject.toml 2026-05-28 09:00:17 +10:00
Robert Wojciechowski d6ce26f143 Merge pull request #84 from robertvoy/fix/issue-83
Fix collector handling of ComfyUI list inputs
2026-05-28 08:59:02 +10:00
Robert Wojciechowski e2125c7529 Fix collector handling of ComfyUI list inputs 2026-05-26 23:02:26 +00:00
Robert Wojciechowski a91f9fb081 Update pyproject.toml 2026-04-08 10:01:49 +10:00
Robert Wojciechowski 6874d2735f Merge branch 'fix/issue-78-queue-status' 2026-04-07 23:41:58 +00:00
Robert Wojciechowski 0e9ee7fbac test: make async helper regression portable 2026-04-07 23:39:18 +00:00
Robert Wojciechowski 5386de10e3 fix: include queue metadata timestamps 2026-04-07 23:34:15 +00:00
Robert Wojciechowski 41f9e44945 fix: return native prompt queue metadata 2026-04-07 23:19:37 +00:00
Robert Wojciechowski 7792d613bb Update FUNDING.yml 2026-04-04 17:42:09 +11:00
Robert Wojciechowski 6e698512b7 Update pyproject.toml 2026-04-04 17:36:24 +11:00
Robert Wojciechowski 79cc9f5ad5 Update registry publish action ref 2026-04-04 06:35:44 +00:00
Robert Wojciechowski 4eba87cba3 Update pyproject.toml 2026-04-04 17:09:55 +11:00
Robert Wojciechowski 27e08de94c Document ComfyUI Desktop support 2026-04-04 06:06:59 +00:00
Robert Wojciechowski c8453a139b Use loopback callback URLs for local workers 2026-04-04 00:15:00 +00:00
Robert Wojciechowski aae831e1e5 Fix stale master callback port selection 2026-04-04 00:05:49 +00:00
Robert Wojciechowski e7ab67733b Add ComfyUI Desktop worker support 2026-04-03 21:54:51 +00:00
Robert Wojciechowski dd55ff740e fix: forward all queuePrompt args through interceptor
The interceptor only captured (number, prompt) and dropped the third
options argument. This caused partialExecutionTargets to be lost,
making ComfyUI execute all output nodes instead of just the selected
one when using Execute Selected Output.

Fixes #76
2026-03-26 22:00:53 +00:00
Robert Wojciechowski a6d0b82d35 Update README.md 2026-03-02 16:10:42 +11:00
Robert Wojciechowski a694272b40 Harden worker probe response validation and add tests 2026-03-01 22:36:57 +00:00
Robert Wojciechowski 16ea22a643 Update pyproject.toml 2026-02-28 15:58:46 +11:00
Robert Wojciechowski 071d1c1d93 Update pyproject.toml 2026-02-28 04:53:20 +00:00
157 changed files with 30829 additions and 741 deletions
-1
View File
@@ -1,4 +1,3 @@
# These are supported funding model platforms
github: robertvoy
buy_me_a_coffee: robertvoy
+1 -1
View File
@@ -19,6 +19,6 @@ jobs:
- name: Check out code
uses: actions/checkout@v4
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@v1
uses: Comfy-Org/publish-node-action@main
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+76 -58
View File
@@ -15,7 +15,7 @@
---
## Key Features
## Key Features
#### Parallel Workflow Processing
- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
@@ -27,32 +27,28 @@
- Intelligent distribution
- Handles single images and videos
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
- Convert any workflow to distributed with 2 nodes
- JSON configuration with UI controls
---
## Current Architecture
- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
- There is **no Distributed Queue node** anymore.
- With `load_balance=true`, orchestration selects one least-busy execution participant:
- If master participation is enabled, master is included as a candidate.
- If master is in orchestrator-only mode, only workers are considered.
---
#### Dynamic Image-List Processing
- Process a whole image list in parallel: the master and every worker take the next image when they free up, and results come back in input order
- Use **Distributed Image List Start / End** around a per-image processing region
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
- Convert any workflow to distributed with 2 nodes
- JSON configuration with UI controls
---
## Worker Types
<img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
ComfyUI Distributed supports three types of workers:
ComfyUI Distributed supports four types of workers:
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
- **Remote Workers** - GPUs on other computers in your network
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
- **Serverless (Runpod)** - Queue-based endpoints for image Collector jobs; no Distributed install or inbound tunnel on the endpoint. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
@@ -61,7 +57,6 @@ ComfyUI Distributed supports three types of workers:
## Requirements
- ComfyUI
> Note: Desktop app not currently supported
- Multiple NVIDIA GPUs
> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
- That's it
@@ -92,19 +87,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
## Workflow Examples
### Basic Parallel Generation
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
### Basic Parallel Generation
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb)
> [Download workflow](/workflows/distributed-txt2img.json)
1. Open your ComfyUI workflow
2. Add **Distributed Seed** → connect to sampler's seed
3. Add **Distributed Collector** → after VAE Decode
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
5. Enable workers in the UI
6. Run the workflow!
2. Add **Distributed Seed** → connect to sampler's seed
3. Add **Distributed Collector** → after VAE Decode
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
5. Enable workers in the UI
6. Run the workflow!
### Parallel WAN Generation
Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
@@ -122,12 +117,12 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
7. Enable workers in the UI
8. Run the workflow!
### Distributed Image Upscaling
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
### Distributed Image Upscaling
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a)
> [Download workflow](/workflows/distributed-upscale.json)
> [Download workflow](/workflows/distributed-upscale.json)
1. Load your image
2. Upscale with ESRGAN or similar
@@ -153,42 +148,44 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
---
## Developer API
## Developer API
Control your distributed cluster programmatically without opening the browser.
* **Endpoint:** `POST /distributed/queue`
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
---
## Distributed Value
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
- The node shows only currently enabled workers.
- If worker enablement changes, worker fields update automatically.
- When disconnected, it resets to default string mode and clears per-worker overrides.
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
---
## Nodes
| Node | Description |
|------|-------------|
| **Distributed Seed** | Generates unique seeds for each worker |
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
| **Distributed Value** | Outputs per-worker override values with fallback to default |
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
| **Image Batch Divider** | Splits image batches for multi-GPU output |
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
---
## Distributed Value
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
- The node shows only currently enabled workers.
- If worker enablement changes, worker fields update automatically.
- When disconnected, it resets to default string mode and clears per-worker overrides.
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
---
## Nodes
| Node | Description |
|------|-------------|
| **Distributed Seed** | Generates unique seeds for each worker |
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
| **Distributed Value** | Outputs per-worker override values with fallback to default |
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
| **Image Batch Divider** | Splits image batches for multi-GPU output |
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
| **Distributed Image List Start** | Turns an image list or batch into one job per image, scheduled to the next available participant; mask and value inputs are shared across images or paired one per image |
| **Distributed Image List End** | Closes the paired region and returns images, masks, and texts as lists in the original input order |
---
@@ -206,7 +203,7 @@ No, it does not speed up the generation of a single image or video. Instead, it
<details>
<summary>Does it work with the ComfyUI desktop app?</summary>
Currently, it is not compatible with the ComfyUI desktop app.
Yes, it does now.
</details>
<details>
@@ -231,6 +228,26 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do
---
## Community projects
Community-maintained add-ons and forks that extend ComfyUI-Distributed.
### [ComfyUI-H3MC-Distributed](https://github.com/ymiyad/ComfyUI-H3MC-Distributed) by ymiyad
Connects H3 Motion Context with ComfyUI-Distributed to generate multiple candidate takes of the next video clip across your workers. Compare the results, pick your preferred take, and use it as context for the next clip.
Each saved candidate video includes its continuation latent, so you don't need a shared folder or separate latent files. This add-on speeds up iteration through parallel candidates; it does not split a single clip across GPUs.
### [ComfyUI-Deadline-Distributed](https://github.com/doubletwisted/ComfyUI-Deadline-Distributed) by doubletwisted
A fork of ComfyUI-Distributed with Thinkbox/AWS Deadline 10 integration for render farms. Claim Deadline workers from the Distributed panel, register them automatically, and collect their results back into your master workflow.
Requires Deadline 10 and the author's ComfyUI-Deadline-Plugin.
See each repository for installation instructions, requirements, and support.
---
## Disclaimer
This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats.
@@ -249,3 +266,4 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
+15 -22
View File
@@ -1,29 +1,22 @@
# Import everything needed from the main module
from .distributed import (
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
)
"""ComfyUI-Distributed's native V3 extension entrypoint."""
from comfy_api.v0_0_2 import ComfyExtension, io
# Import utilities
from .utils.config import ensure_config_exists, CONFIG_FILE
from .utils.logging import debug_log
from .nodes.v3 import NODES
from .runtime.bootstrap import initialize
# Import distributed upscale nodes
from .nodes.distributed_upscale import (
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
)
WEB_DIRECTORY = './web'
WEB_DIRECTORY = "./web"
ensure_config_exists()
class DistributedExtension(ComfyExtension):
async def on_load(self) -> None:
initialize()
# Merge node mappings
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return list(NODES)
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
debug_log("Loaded Distributed nodes.")
debug_log(f"Config file: {CONFIG_FILE}")
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
async def comfy_entrypoint() -> DistributedExtension:
return DistributedExtension()
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
+2
View File
@@ -1,5 +1,7 @@
from . import config_routes # noqa: F401
from . import serverless_routes # noqa: F401
from . import tunnel_routes # noqa: F401
from . import worker_routes # noqa: F401
from . import job_routes # noqa: F401
from . import usdu_routes # noqa: F401
from . import image_list_routes # noqa: F401
+382 -69
View File
@@ -1,5 +1,6 @@
import json
from contextlib import asynccontextmanager
from typing import Any
from aiohttp import web
import server
@@ -24,6 +25,12 @@ except ImportError:
save_config(config)
from ..utils.logging import debug_log, log
from ..utils.network import handle_api_error, normalize_host
from ..utils.secret_fields import is_secret_field_name
from ..utils.serverless_config import (
ServerlessSecretChanges,
public_config,
validate_serverless_worker,
)
def _positive_int(value):
@@ -86,22 +93,169 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
target[key] = normalizer(value) if (normalizer and value is not None) else value
_SECRET_FIELDS = {"api_key", "clear_api_key"}
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
_SERVERLESS_ONLY_FIELDS = {
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
"media_profile",
}
_SERVERLESS_CONFIG_FIELDS = {
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
}
class _ActivePoolCheckError(RuntimeError):
pass
def _error_response(message, status):
return web.json_response({"status": "error", "message": message}, status=status)
def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0):
if isinstance(value, dict):
for key, item in value.items():
field_name = key.casefold() if isinstance(key, str) else ""
if is_secret_field_name(field_name):
if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS):
return True
if _contains_secret_fields(
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
):
return True
elif isinstance(value, (list, tuple)):
return any(
_contains_secret_fields(
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
)
for item in value
)
return False
def _serverless_secret_store():
from .orchestration.serverless_runtime import get_secret_store
return get_secret_store()
def _serverless_pool_is_active(worker_id):
from .orchestration.serverless_runtime import get_manager
return get_manager().pool_is_active(worker_id)
def _ensure_pool_idle(worker_id):
try:
active = _serverless_pool_is_active(worker_id)
except Exception:
raise _ActivePoolCheckError from None
if active:
raise _ActivePoolCheckError
def _worker_by_id(workers, worker_id):
for worker in workers:
if isinstance(worker, dict) and worker.get("id") == worker_id:
return worker
return None
def _worker_type(worker):
return worker.get("type", "local") if isinstance(worker, dict) else None
def _strip_sensitive_fields(value: Any) -> Any:
if isinstance(value, dict):
return {
key: _strip_sensitive_fields(item)
for key, item in value.items()
if not is_secret_field_name(key)
}
if isinstance(value, list):
return [_strip_sensitive_fields(item) for item in value]
return value
def _clean_serverless_worker(worker):
clean = _strip_sensitive_fields(dict(worker))
for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}:
clean.pop(field, None)
return clean
def _validate_bulk_workers(workers):
validated = []
for worker in workers:
if isinstance(worker, dict) and worker.get("type") == "serverless":
clean = _clean_serverless_worker(worker)
validated.append(validate_serverless_worker(clean))
else:
validated.append(worker)
return validated
def _clean_bulk_type_conversions(old_workers, new_workers):
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
for index, worker in enumerate(new_workers):
if not isinstance(worker, dict) or worker.get("id") is None:
continue
old = old_by_id.get(str(worker["id"]))
if old is None or old.get("type") != "serverless":
continue
if worker.get("type") == "serverless":
# Older clients do not know this field. Only an explicit selection
# may change the endpoint contract, including on bulk saves.
if "media_profile" not in worker and "media_profile" in old:
worker["media_profile"] = old["media_profile"]
continue
converted = dict(worker)
for field in _SERVERLESS_ONLY_FIELDS:
converted.pop(field, None)
new_workers[index] = converted
return new_workers
def _bulk_serverless_changes(old_workers, new_workers):
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
secret_ids = set()
for worker_id, old in old_by_id.items():
new = new_by_id.get(worker_id)
if old.get("type") == "serverless":
if new is None or new.get("type") != "serverless":
_ensure_pool_idle(worker_id)
secret_ids.add(worker_id)
elif (
old.get("endpoint_id") != new.get("endpoint_id")
or old.get("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
):
_ensure_pool_idle(worker_id)
elif new is not None and new.get("type") == "serverless":
secret_ids.add(worker_id)
return secret_ids
@server.PromptServer.instance.routes.get("/distributed/config")
async def get_config_endpoint(request):
config = load_config()
return web.json_response(config)
try:
config = load_config()
return web.json_response(public_config(config, _serverless_secret_store()))
except Exception:
return _error_response("Unable to read config", 500)
@server.PromptServer.instance.routes.post("/distributed/config")
async def update_config_endpoint(request):
"""Bulk config update with schema validation."""
"""Bulk config update with schema validation and credential redaction."""
try:
data = await request.json()
except Exception as e:
return await handle_api_error(request, f"Invalid JSON payload: {e}", 400)
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return await handle_api_error(request, "Config payload must be an object", 400)
return _error_response("Config payload must be an object", 400)
if _contains_secret_fields(data):
return _error_response("API keys cannot be set through bulk config", 400)
validated_settings = {}
validated_root = {}
@@ -118,9 +272,15 @@ async def update_config_endpoint(request):
continue
if validator and not validator(value):
errors.append(f"{key}: value {value!r} failed validation")
errors.append(f"{key}: failed validation")
continue
if key == "workers":
try:
value = _validate_bulk_workers(value)
except (TypeError, ValueError):
return _error_response("Invalid serverless worker configuration", 400)
if key in _SETTINGS_FIELDS:
validated_settings[key] = value
else:
@@ -133,15 +293,41 @@ async def update_config_endpoint(request):
"message": "; ".join(errors),
}, status=400)
store = None
response_config = None
try:
async with config_transaction() as config:
settings = config.setdefault("settings", {})
settings.update(validated_settings)
for key, value in validated_root.items():
config[key] = value
return web.json_response({"status": "success", "config": config})
except Exception as e:
return await handle_api_error(request, e)
with ServerlessSecretChanges() as secret_changes:
async with config_transaction() as config:
if "workers" in validated_root:
old_workers = config.get("workers", [])
validated_root["workers"] = _clean_bulk_type_conversions(
old_workers, validated_root["workers"]
)
changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"])
for worker in validated_root["workers"]:
if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False):
store = store or _serverless_secret_store()
if worker["id"] in changed_ids or not store.has(worker["id"]):
return _error_response("Save an API key before enabling this worker", 400)
if changed_ids:
store = store or _serverless_secret_store()
for worker_id in sorted(changed_ids):
prior_secret = secret_changes.capture(store, worker_id)
if prior_secret is not None:
secret_changes.delete(store, worker_id)
settings = config.setdefault("settings", {})
settings.update(validated_settings)
for key, value in validated_root.items():
config[key] = value
response_config = public_config(config, store or _serverless_secret_store())
response = web.json_response({"status": "success", "config": response_config})
secret_changes.commit()
return response
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
except Exception:
return _error_response("Unable to update config", 500)
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
@@ -164,79 +350,206 @@ async def queue_status_endpoint(request):
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
async def update_worker_endpoint(request):
return await _run_with_serverless_secret_changes(_update_worker_endpoint, request)
async def _run_with_serverless_secret_changes(endpoint, request):
with ServerlessSecretChanges() as secret_changes:
response = await endpoint(request, secret_changes)
if response.status < 400:
secret_changes.commit()
return response
async def _update_worker_endpoint(request, secret_changes):
store = None
worker_id = None
try:
data = await request.json()
try:
data = await request.json()
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return _error_response("Worker payload must be an object", 400)
worker_id = data.get("worker_id")
if worker_id is None:
return await handle_api_error(request, "Missing worker_id", 400)
return _error_response("Missing worker_id", 400)
if _contains_secret_fields(data, allow_root_serverless_secrets=True):
return _error_response("Invalid worker credentials", 400)
result = None
async with config_transaction() as config:
worker_found = False
workers = config.setdefault("workers", [])
existing = _worker_by_id(workers, worker_id)
current_type = _worker_type(existing)
target_type = data.get("type", current_type or "local")
for worker in workers:
if worker["id"] == worker_id:
_apply_field_patch(worker, data, _WORKER_FIELDS)
worker_found = True
break
if target_type == "serverless":
unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS
if unknown:
return _error_response("Unknown serverless worker field", 400)
if not isinstance(worker_id, str) or not worker_id.strip():
return _error_response("worker_id must be a non-empty string", 400)
if "api_key" in data and not isinstance(data["api_key"], str):
return _error_response("Invalid worker credentials", 400)
if "clear_api_key" in data and type(data["clear_api_key"]) is not bool:
return _error_response("Invalid worker credentials", 400)
supplied_key = data.get("api_key", "")
clear_key = data.get("clear_api_key", False)
if supplied_key and clear_key:
return _error_response("Choose either api_key or clear_api_key", 400)
if "enabled" in data and type(data["enabled"]) is not bool:
return _error_response("enabled must be a boolean", 400)
if "name" in data and not isinstance(data["name"], str):
return _error_response("name must be a string", 400)
if not worker_found:
# If worker not found and all required fields are provided, create new worker
if all(key in data for key in ["name", "port", "cuda_device"]):
new_worker = {
"id": worker_id,
"name": data["name"],
"host": normalize_host(data.get("host", "localhost")),
"port": data["port"],
"cuda_device": data["cuda_device"],
"enabled": data.get("enabled", False),
"extra_args": data.get("extra_args", ""),
"type": data.get("type", "local")
}
workers.append(new_worker)
candidate = dict(existing or {})
candidate["id"] = worker_id
candidate["type"] = "serverless"
candidate.setdefault("provider", "runpod")
candidate.setdefault("enabled", False)
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
if field in data:
candidate[field] = data[field]
candidate = _clean_serverless_worker(candidate)
validated = validate_serverless_worker(candidate)
if "enabled" in validated and type(validated["enabled"]) is not bool:
return _error_response("enabled must be a boolean", 400)
if "name" in validated and not isinstance(validated["name"], str):
return _error_response("name must be a string", 400)
store = _serverless_secret_store()
prior_secret = secret_changes.capture(store, worker_id)
converting_to_serverless = current_type != "serverless"
next_secret = None if converting_to_serverless else prior_secret
if supplied_key:
next_secret = supplied_key
elif clear_key:
next_secret = None
endpoint_changed = (
current_type == "serverless"
and existing.get("endpoint_id") != validated.get("endpoint_id")
)
media_profile_changed = (
current_type == "serverless"
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
)
if current_type == "serverless" and (
endpoint_changed or media_profile_changed or next_secret != prior_secret
):
_ensure_pool_idle(worker_id)
if validated.get("enabled", False) and next_secret is None:
return _error_response("Save an API key before enabling this worker; turn the worker off before clearing its key", 400)
if next_secret != prior_secret:
if next_secret is None:
secret_changes.delete(store, worker_id)
else:
secret_changes.set(store, worker_id, next_secret)
if existing is None:
workers.append(validated)
else:
return await handle_api_error(
request,
f"Worker {worker_id} not found and missing required fields for creation",
404,
)
existing.clear()
existing.update(validated)
result = {"status": "success", "api_key_configured": next_secret is not None}
elif existing is not None:
cleaned_existing = _strip_sensitive_fields(existing)
existing.clear()
existing.update(cleaned_existing)
if current_type == "serverless":
if _contains_secret_fields(data):
return _error_response("Invalid worker credentials", 400)
_ensure_pool_idle(worker_id)
store = _serverless_secret_store()
prior_secret = secret_changes.capture(store, worker_id)
if prior_secret is not None:
secret_changes.delete(store, worker_id)
for field in _SERVERLESS_ONLY_FIELDS:
existing.pop(field, None)
elif _contains_secret_fields(data):
return _error_response("API keys are only accepted for serverless workers", 400)
_apply_field_patch(existing, data, _WORKER_FIELDS)
result = {"status": "success"}
elif all(key in data for key in ["name", "port", "cuda_device"]):
if _contains_secret_fields(data):
return _error_response("API keys are only accepted for serverless workers", 400)
new_worker = {
"id": worker_id,
"name": data["name"],
"host": normalize_host(data.get("host", "localhost")),
"port": data["port"],
"cuda_device": data["cuda_device"],
"enabled": data.get("enabled", False),
"extra_args": data.get("extra_args", ""),
"type": data.get("type", "local")
}
workers.append(new_worker)
result = {"status": "success"}
else:
return _error_response(
"Worker not found and missing required fields for creation", 404
)
return web.json_response(result)
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
except ValueError:
return _error_response("Invalid worker configuration", 400)
except Exception:
return _error_response("Unable to update worker configuration", 500)
return web.json_response({"status": "success"})
except Exception as e:
return await handle_api_error(request, e, 400)
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
async def delete_worker_endpoint(request):
return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
async def _delete_worker_endpoint(request, secret_changes):
store = None
worker_id = None
try:
data = await request.json()
try:
data = await request.json()
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return _error_response("Worker payload must be an object", 400)
worker_id = data.get("worker_id")
if worker_id is None:
return await handle_api_error(request, "Missing worker_id", 400)
return _error_response("Missing worker_id", 400)
result = None
async with config_transaction() as config:
workers = config.get("workers", [])
# Find and remove the worker
worker_index = -1
for i, worker in enumerate(workers):
if worker["id"] == worker_id:
worker_index = i
break
if worker_index == -1:
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
# Remove the worker
removed_worker = workers.pop(worker_index)
return web.json_response({
removed = _worker_by_id(workers, worker_id)
if removed is None:
return _error_response("Worker not found", 404)
if _worker_type(removed) == "serverless":
_ensure_pool_idle(worker_id)
store = _serverless_secret_store()
secret_worker_id = str(worker_id)
prior_secret = secret_changes.capture(store, secret_worker_id)
if prior_secret is not None:
secret_changes.delete(store, secret_worker_id)
workers.remove(removed)
result = {
"status": "success",
"message": f"Worker {removed_worker.get('name', worker_id)} deleted"
})
except Exception as e:
return await handle_api_error(request, e, 400)
"message": "Worker deleted" if _worker_type(removed) == "serverless"
else f"Worker {removed.get('name', worker_id)} deleted",
}
return web.json_response(result)
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
except Exception:
return _error_response("Unable to delete worker", 500)
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
async def update_setting_endpoint(request):
+92
View File
@@ -0,0 +1,92 @@
"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
import json
from aiohttp import web
import server
from ..image_list.runtime import (
IMAGE_TRANSFER_FORMATS,
MAX_REQUEST_BYTES,
PROTOCOL_VERSION,
decode_result,
delete_worker_item,
set_master_result,
store_worker_item,
)
prompt_server = server.PromptServer.instance
async def _json_body(request):
length = request.content_length
if length is not None and length > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
raw = bytearray()
async for chunk in request.content.iter_chunked(64 * 1024):
raw.extend(chunk)
if len(raw) > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
try:
return json.loads(raw.decode("utf-8"))
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
@prompt_server.routes.get("/distributed/image-list/capabilities")
async def image_list_capabilities_endpoint(_request):
return web.json_response({
"supported": True,
"execution_targets": True,
"protocol_version": PROTOCOL_VERSION,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
"transport": "typed-float-tensor-json-v1",
"image_transfer_formats": list(IMAGE_TRANSFER_FORMATS),
})
@prompt_server.routes.post("/distributed/image-list/item")
async def image_list_item_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list item object")
context_id = data.get("context_id")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
item = data.get("item")
store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"]})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
@prompt_server.routes.delete("/distributed/image-list/item/{token}")
async def image_list_delete_item_endpoint(request):
token = request.match_info.get("token", "")
return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
@prompt_server.routes.post("/distributed/image-list/result")
async def image_list_result_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list result object")
context_id = data.get("context_id")
token = data.get("token")
item_index = data.get("item_index")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
if not isinstance(token, str) or not token:
raise ValueError("token must be a non-empty string")
if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
raise ValueError("item_index must be a non-negative integer")
result = decode_result(data.get("result"))
set_master_result(context_id, item_index, token, result)
return web.json_response({"status": "accepted", "item_index": item_index})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
+167 -29
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):
@@ -217,7 +226,7 @@ async def distributed_queue_endpoint(request):
return await handle_api_error(request, exc, 400)
try:
prompt_id, worker_count = await orchestrate_distributed_execution(
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
payload.prompt,
payload.workflow_meta,
payload.client_id,
@@ -227,6 +236,8 @@ async def distributed_queue_endpoint(request):
)
return web.json_response({
"prompt_id": prompt_id,
"number": prompt_number,
"node_errors": node_errors,
"worker_count": worker_count,
"auto_prepare_supported": True,
})
@@ -268,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:
@@ -293,42 +432,41 @@ async def job_complete_endpoint(request):
errors.append("worker_id: expected non-empty string")
if not isinstance(batch_idx, int) or batch_idx < 0:
errors.append("batch_idx: expected non-negative integer")
if not isinstance(image_payload, str) or not image_payload.strip():
image_provided = image_payload is not None
audio_provided = audio_payload is not None
if not image_provided and not audio_provided:
errors.append("expected at least one of image or audio")
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
errors.append("image: expected non-empty base64 PNG string")
if audio_payload is not None and not isinstance(audio_payload, dict):
if audio_provided and not isinstance(audio_payload, dict):
errors.append("audio: expected object when provided")
if not isinstance(is_last, bool):
errors.append("is_last: expected boolean")
if errors:
return await handle_api_error(request, errors, 400)
tensor = _decode_canonical_png_tensor(image_payload)
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
try:
tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
except ValueError as exc:
return await handle_api_error(request, exc, 400)
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}, "
+179 -22
View File
@@ -6,6 +6,9 @@ import aiohttp
from ...utils.logging import debug_log, log
from ...utils.network import build_worker_url, get_client_session, probe_worker
# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module
# stays importable without the image-list runtime; a test pins them equal).
IMAGE_LIST_PROTOCOL_VERSION = 2
try:
from ...utils.trace_logger import trace_debug, trace_info
except ImportError:
@@ -55,6 +58,11 @@ async def worker_ws_is_active(worker):
async def _probe_worker_active(worker, use_websocket, semaphore):
async with semaphore:
if worker.get("type") == "serverless":
from .serverless_runtime import get_manager
# Authentication/configuration failures must surface, not silently skip paid work.
status = await get_manager().health(worker)
return worker, status["ok"]
is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
return worker, is_active
@@ -69,30 +77,49 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
"workflow": payload.get("workflow"),
"client_id": client_id,
}
if "partial_execution_targets" in payload:
ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
ws_url = f"{ws_url}/distributed/worker_ws"
session = await get_client_session()
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
ws = None
async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
await ws.send_json(ws_payload)
async for msg in ws:
if msg.type == aiohttp.WSMsgType.TEXT:
data = json.loads(msg.data or "{}")
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
if data.get("ok"):
return
error_text = data.get("error") or "Worker rejected websocket dispatch."
validation_error = data.get("validation_error")
node_errors = data.get("node_errors")
if validation_error:
error_text = f"{error_text} | validation_error={validation_error}"
if node_errors:
error_text = f"{error_text} | node_errors={node_errors}"
raise RuntimeError(error_text)
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
try:
async with asyncio.timeout_at(deadline):
ws = await session.ws_connect(ws_url, heartbeat=20)
await ws.send_json(ws_payload)
while True:
msg = await ws.receive()
if msg.type == aiohttp.WSMsgType.TEXT:
data = json.loads(msg.data or "{}")
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
if data.get("ok"):
return data.get("prompt_id")
error_text = data.get("error") or "Worker rejected websocket dispatch."
validation_error = data.get("validation_error")
node_errors = data.get("node_errors")
if validation_error:
error_text = f"{error_text} | validation_error={validation_error}"
if node_errors:
error_text = f"{error_text} | node_errors={node_errors}"
raise RuntimeError(error_text)
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
except asyncio.TimeoutError as exc:
raise RuntimeError(
"Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
) from exc
finally:
if ws is not None:
# Closing must not discard an already acknowledged prompt ID or
# replace a definite rejection with an ambiguous timeout.
try:
async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
await ws.close()
except (asyncio.TimeoutError, aiohttp.ClientError):
pass
async def dispatch_worker_prompt(
@@ -102,11 +129,18 @@ async def dispatch_worker_prompt(
client_id=None,
use_websocket=False,
trace_execution_id=None,
return_prompt_id=False,
*,
partial_execution_targets=None,
):
"""Send the prepared prompt to a worker ComfyUI instance."""
worker_url = build_worker_url(worker)
url = build_worker_url(worker, "/prompt")
payload = {"prompt": prompt_obj}
target_payload = {}
if partial_execution_targets is not None:
target_payload["partial_execution_targets"] = list(partial_execution_targets)
payload.update(target_payload)
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
@@ -115,15 +149,18 @@ async def dispatch_worker_prompt(
if use_websocket:
try:
await _dispatch_via_websocket(
prompt_id = await _dispatch_via_websocket(
worker_url,
{
"prompt": prompt_obj,
"workflow": workflow_meta,
**target_payload,
},
client_id,
)
return
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
return prompt_id if return_prompt_id else None
except Exception as exc:
worker_id = worker.get("id")
if trace_execution_id:
@@ -139,6 +176,122 @@ async def dispatch_worker_prompt(
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError("Worker /prompt response omitted prompt_id")
return prompt_id if return_prompt_id else None
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
serverless provider. Unsupported worker releases fail before master inputs run.
"""
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)
async def inspect(worker):
worker_id = str(worker.get("id") or "unknown")
check = "capability"
try:
async with semaphore:
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(
f"Image-list worker {worker_id} does not expose the capability endpoint "
f"(HTTP {response.status}); upgrade that worker before running this feature"
)
capability = await response.json()
if not isinstance(capability, dict):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
supported = capability.get("supported")
protocol_version = capability.get("protocol_version")
nodes = capability.get("node_ids")
if (
not isinstance(supported, bool)
or type(protocol_version) is not int
or not isinstance(nodes, list)
or any(not isinstance(node_id, str) for node_id in nodes)
):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
if (
supported is not True
or protocol_version != IMAGE_LIST_PROTOCOL_VERSION
or not required_nodes.issubset(nodes)
):
raise RuntimeError(
f"Image-list worker {worker_id} is incompatible: requires protocol "
f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes "
f"{sorted(required_nodes)}; no work was submitted"
)
if capability.get("execution_targets") is not True:
raise RuntimeError(
f"Image-list worker {worker_id} lacks explicit execution_targets support; "
"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:
if response.status != 200:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
queue = await response.json()
if (
not isinstance(queue, dict)
or not isinstance(queue.get("queue_running"), list)
or not isinstance(queue.get("queue_pending"), list)
):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
idle = not queue["queue_running"] and not queue["queue_pending"]
return worker_id, idle
except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
log(
f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
f"skipping: {exc}"
)
return worker_id, None
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
idle_workers = {
worker_id: True
for worker_id, idle in statuses
if idle is True
}
if require_idle and not idle_workers:
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
return idle_workers
async def select_active_workers(
@@ -203,6 +356,10 @@ def _extract_queue_remaining(payload):
async def _probe_worker_queue(worker, semaphore, probe_timeout):
async with semaphore:
if worker.get("type") == "serverless":
from .serverless_runtime import get_manager
status = await get_manager().health(worker)
return {"worker": worker, "queue_remaining": status["queueRemaining"]}
worker_url = build_worker_url(worker)
payload = await probe_worker(worker_url, timeout=probe_timeout)
if payload is None:
+35
View File
@@ -29,6 +29,11 @@ def _normalize_media_reference(value):
return None
cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
if MEDIA_FILE_RE.search(cleaned):
# Output/temp annotations select a different master-side media root.
# Preserve them through discovery/loading and as distinct rewrite keys.
annotation = re.search(r"\s*\[(output|temp)\]$", value)
if annotation:
return f"{cleaned} [{annotation.group(1)}]"
return cleaned
return None
@@ -143,10 +148,40 @@ 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()
normalized = filename.replace("\\", "/")
annotation = re.search(r"\s*\[(output|temp)\]$", normalized)
if annotation:
# Worker uploads live under input/. Do not collide across master roots,
# or overwrite a different reference with the same basename.
basename = normalized[:annotation.start()].rstrip().rsplit("/", 1)[-1]
normalized = f"distributed-media/{annotation.group(1)}/{file_hash}/{basename}"
check_url = build_worker_url(worker, "/distributed/check_file")
try:
+351 -16
View File
@@ -125,8 +125,222 @@ def _find_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."""
_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
"PrimitiveBoolean",
"PrimitiveFloat",
"PrimitiveInt",
"PrimitiveNode",
"PrimitiveString",
}
_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
"LoadImage",
}
_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
# only after its connected inputs recursively prove safe.
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
"CreateList": ("inputs.",),
}
# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
def _get_delegate_master_node_class_mappings():
"""Return ComfyUI node-class mappings when available."""
if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
try:
import nodes as comfy_nodes # type: ignore
except Exception: # pragma: no cover - depends on ComfyUI runtime imports
return {}
return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
def _get_delegate_master_node_class(class_type):
mappings = _get_delegate_master_node_class_mappings()
return mappings.get(class_type) if isinstance(mappings, dict) else None
def _normalize_delegate_master_return_type(return_type):
if return_type is None:
return ""
return str(return_type).strip().upper()
def _delegate_master_type_is_safe_scalar(type_name):
return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
def _delegate_master_type_is_safe_config(type_name):
return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
def _delegate_master_output_is_safe_scalar(class_type, output_index):
"""Return True when a registered node output is lightweight config data."""
if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
return True
node_class = _get_delegate_master_node_class(class_type)
return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
try:
output_type = return_types[int(output_index)]
except (IndexError, TypeError, ValueError):
return False
return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
def _get_delegate_master_input_types(class_type):
node_class = _get_delegate_master_node_class(class_type)
input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
if callable(input_types):
try:
input_types = input_types()
except TypeError:
return {}
return input_types if isinstance(input_types, dict) else {}
def _normalize_delegate_master_input_type(input_spec):
if isinstance(input_spec, (list, tuple)) and input_spec:
return _normalize_delegate_master_return_type(input_spec[0])
return _normalize_delegate_master_return_type(input_spec)
def _delegate_master_input_is_safe_scalar(class_type, input_name):
"""Return True when a registered downstream input expects config data."""
for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
if input_name.startswith(prefix):
return True
input_types = _get_delegate_master_input_types(class_type)
for section_name in ("required", "optional"):
section = input_types.get(section_name, {})
if isinstance(section, dict) and input_name in section:
input_type = _normalize_delegate_master_input_type(section[input_name])
return _delegate_master_type_is_safe_config(input_type)
return False
def _is_delegate_master_always_retained_upstream_node(node):
if not isinstance(node, dict):
return False
class_type = node.get("class_type")
return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
def _is_delegate_master_retained_upstream_node(node, output_index=0):
"""Return True for lightweight upstream nodes safe to keep on the master."""
if not isinstance(node, dict):
return False
class_type = node.get("class_type")
if not isinstance(class_type, str):
return False
return (
class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
or class_type.startswith("Primitive")
or _delegate_master_output_is_safe_scalar(class_type, output_index)
)
def _collect_delegate_master_retained_upstream_branch(
prompt_obj,
node_id,
output_index,
memo,
visiting,
):
"""Return safe retained branch nodes, or None when the branch is not safe."""
node_id = str(node_id)
cache_key = (node_id, output_index)
if cache_key in memo:
cached = memo[cache_key]
return None if cached is None else set(cached)
if cache_key in visiting:
memo[cache_key] = None
return None
node = prompt_obj.get(node_id)
if not _is_delegate_master_retained_upstream_node(node, output_index):
memo[cache_key] = None
return None
visiting.add(cache_key)
retained = {node_id}
inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
class_type = node.get("class_type") if isinstance(node, dict) else None
for input_name, value in inputs.items():
if not (isinstance(value, list) and len(value) == 2):
continue
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
visiting.remove(cache_key)
memo[cache_key] = None
return None
source_id = str(value[0])
branch = _collect_delegate_master_retained_upstream_branch(
prompt_obj,
source_id,
value[1],
memo,
visiting,
)
if branch is None:
visiting.remove(cache_key)
memo[cache_key] = None
return None
retained.update(branch)
visiting.remove(cache_key)
memo[cache_key] = frozenset(retained)
return retained
def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
"""Return lightweight upstream nodes needed by kept delegate-master nodes."""
connected = set()
memo = {}
for node_id in start_ids:
node = prompt_obj.get(str(node_id)) or {}
inputs = node.get("inputs", {})
class_type = node.get("class_type") if isinstance(node, dict) else None
for input_name, value in inputs.items():
if not (isinstance(value, list) and len(value) == 2):
continue
source_node = prompt_obj.get(str(value[0]))
if _is_delegate_master_always_retained_upstream_node(source_node):
connected.add(str(value[0]))
continue
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
continue
branch = _collect_delegate_master_retained_upstream_branch(
prompt_obj,
value[0],
value[1],
memo,
set(),
)
if branch is not None:
connected.update(branch)
return connected
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
@@ -148,25 +362,65 @@ def prune_prompt_for_worker(prompt_obj):
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
if has_removed_downstream:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {
"images": [dist_id, 0],
},
"class_type": "PreviewImage",
"_meta": {
"title": "Preview Image (auto-added)",
},
}
original_node = prompt_obj.get(str(dist_id), {})
class_type = original_node.get("class_type")
inputs = original_node.get("inputs", {})
image_connected = class_type != "DistributedCollector" or (
isinstance(inputs.get("images"), list)
and len(inputs["images"]) == 2
)
audio_connected = (
class_type == "DistributedCollector"
and isinstance(inputs.get("audio"), list)
and len(inputs["audio"]) == 2
)
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]},
"class_type": "PreviewImage",
"_meta": {"title": "Preview Image (auto-added)"},
}
elif audio_connected:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {"audio": [dist_id, 1]},
"class_type": "PreviewAudio",
"_meta": {"title": "Preview Audio (auto-added)"},
}
return pruned_prompt
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
# Typed collectors need no fabricated IMAGE placeholder. Sever all media
# links, including STRING scalar branches, before retained-config analysis.
# Their expected_media control was frozen by participant overrides.
if any((prompt_obj.get(cid) or {}).get("class_type") == "DistributedMediaCollector"
for cid in collector_ids):
from ...serverless_bridge.contract import SOCKET_KINDS
prompt_obj = json.loads(json.dumps(prompt_obj))
for cid in collector_ids:
node = prompt_obj.get(cid) or {}
if node.get("class_type") == "DistributedMediaCollector":
inputs = node.setdefault("inputs", {})
inputs.setdefault("expected_media", json.dumps([name for name in SOCKET_KINDS if name in inputs]))
for name in SOCKET_KINDS:
inputs.pop(name, None)
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
nodes_to_keep = set(collector_ids)
nodes_to_keep.update(downstream)
nodes_to_keep.update(
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
)
pruned_prompt = {}
for node_id in nodes_to_keep:
@@ -194,6 +448,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
collector_entry = pruned_prompt.get(collector_id)
if not collector_entry:
continue
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
original_images = original_inputs.get("images")
if not (isinstance(original_images, list) and len(original_images) == 2):
continue
placeholder_id = next_id()
pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage",
@@ -214,11 +472,56 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
return pruned_prompt
def _resolve_direct_workers_for_image_list(config, requested_ids):
"""Resolve enabled workers by default, or exact explicitly selected direct workers."""
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
configured_workers = (config or {}).get("workers", [])
configured_ids = {
str(worker.get("id") or "").strip()
for worker in configured_workers
if str(worker.get("id") or "").strip()
}
if requested is not None:
unknown_ids = requested - configured_ids
if unknown_ids:
raise RuntimeError(
"Unknown explicitly selected image-list worker ID(s): "
+ ", ".join(sorted(unknown_ids))
)
direct = []
for worker in configured_workers:
worker_id = str(worker.get("id") or "").strip()
if not worker_id:
continue
if requested is None:
if not worker.get("enabled", False):
continue
elif worker_id not in requested:
continue
if str(worker.get("type") or "local").strip().lower() == "serverless":
continue
try:
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
direct.append({
"id": worker_id,
"name": str(worker.get("name") or worker_id),
"host": worker.get("host"),
"port": port,
"type": worker.get("type", "local"),
})
return direct
def generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
"UltimateSDUpscaleDistributed"
distributed_nodes = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("DistributedMediaCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+ prompt_index.nodes_for_class("DistributedImageListStart")
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -249,8 +552,10 @@ def _override_collector_nodes(
enabled_json,
delegate_master,
):
"""Configure DistributedCollector nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
"""Configure legacy and typed collectors without changing their media."""
collector_ids = (prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("DistributedMediaCollector"))
for node_id in collector_ids:
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
@@ -263,6 +568,9 @@ def _override_collector_nodes(
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
if node.get("class_type") == "DistributedMediaCollector":
from ...serverless_bridge.contract import SOCKET_KINDS
inputs["expected_media"] = json.dumps([socket for socket in SOCKET_KINDS if socket in inputs])
if is_master:
inputs["delegate_only"] = bool(delegate_master)
inputs.pop("master_url", None)
@@ -313,6 +621,22 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
job_id_map, master_url, delegate_master, workflow_meta=None):
"""Attach protocol context to Start without changing its visible V3 sockets."""
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
inputs["image_list_master_participates"] = participant_id == "master" and not delegate_master
inputs["image_list_master_callback"] = str(master_url or "")
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
if workflow_meta is not None:
inputs["image_list_workflow_meta"] = workflow_meta
def apply_participant_overrides(
prompt_copy,
participant_id,
@@ -321,6 +645,7 @@ def apply_participant_overrides(
master_url,
delegate_master,
prompt_index,
workflow_meta=None,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
is_master = participant_id == "master"
@@ -348,5 +673,15 @@ def apply_participant_overrides(
master_url,
enabled_json,
)
_override_image_list_nodes(
prompt_copy,
prompt_index,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
delegate_master,
workflow_meta=workflow_meta,
)
return prompt_copy
+127
View File
@@ -0,0 +1,127 @@
"""Small Runpod queue API client. Never retries a potentially billable POST."""
import asyncio
import json
import re
import aiohttp
API_BASE = "https://api.runpod.ai/v2"
MAX_REQUEST_BYTES = 10_000_000
_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z")
class RunpodError(RuntimeError):
"""Sanitized provider failure safe to display in the panel."""
class AmbiguousSubmission(RunpodError):
"""Submission may have been accepted: never retry it automatically."""
def serialize_request(workflow, images, timeout_seconds, *, media=None):
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
payload = {
"input": {"workflow": workflow, "images": images},
"policy": {
"executionTimeout": timeout_seconds * 1000,
"ttl": timeout_seconds * 1000,
},
}
if media is not None:
from ...serverless_bridge.handler import _validate_request
payload["input"]["media"] = media
_validate_request(payload["input"])
body = json.dumps(
payload,
allow_nan=False,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
size = len(body)
if size > MAX_REQUEST_BYTES:
raise ValueError(
f"Serialized serverless request body is {size:,} bytes; Runpod allows at most "
f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw "
"image bytes; reduce image or workflow size."
)
return body
def checked_id(value):
if not isinstance(value, str) or not _ID.fullmatch(value):
raise ValueError("Invalid Runpod endpoint or job ID")
return value
class RunpodClient:
def __init__(self, session, endpoint_id, api_key):
self.session = session
self.endpoint_id = checked_id(endpoint_id)
if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"):
raise ValueError("A valid Runpod API key is required")
self._api_key = api_key
async def _request(self, method, path, payload=None, request_body=None):
# Base origin is code-owned, never a user-supplied URL. Do not follow redirects
# with credentials, or include response bodies/transport errors in exceptions.
url = f"{API_BASE}/{self.endpoint_id}/{path}"
attempts = 3 if method == "GET" else 1
for attempt in range(attempts):
try:
request_kwargs = {
"headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"},
"timeout": aiohttp.ClientTimeout(total=30),
"allow_redirects": False,
}
if request_body is None:
request_kwargs["json"] = payload
else:
request_kwargs["data"] = request_body
async with self.session.request(
method, url,
**request_kwargs,
) as response:
status = response.status
if status != 200:
if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts:
await asyncio.sleep(0.5 * (2 ** attempt))
continue
if method == "POST" and path == "run" and status >= 500:
raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.")
raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})")
try:
result = await response.json()
except (ValueError, aiohttp.ContentTypeError):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None
raise RunpodError("Runpod returned an invalid JSON response") from None
if not isinstance(result, dict):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.")
raise RunpodError("Runpod returned an invalid response")
return result
except (asyncio.TimeoutError, aiohttp.ClientError, OSError):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None
if method == "GET" and attempt + 1 < attempts:
await asyncio.sleep(0.5 * (2 ** attempt))
continue
raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None
raise RunpodError("Runpod request failed")
async def health(self):
return await self._request("GET", "health")
async def submit(self, workflow, images, timeout_seconds, *, media=None):
body = serialize_request(workflow, images, timeout_seconds, media=media)
result = await self._request("POST", "run", request_body=body)
try:
return checked_id(result.get("id"))
except ValueError:
raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None
async def status(self, job_id):
return await self._request("GET", f"status/{checked_id(job_id)}")
async def cancel(self, job_id):
return await self._request("POST", f"cancel/{checked_id(job_id)}")
+57
View File
@@ -0,0 +1,57 @@
"""Private, atomic job-ID journal. Contains identifiers/status only, never payloads."""
import json
import os
from pathlib import Path
import tempfile
import threading
from .serverless_state import TERMINAL
FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"}
class JobLedger:
def __init__(self, path):
self.path = Path(path)
self._lock = threading.RLock()
def _read(self):
if self.path.is_symlink():
raise RuntimeError("Unsafe serverless job journal path")
if not self.path.exists():
return {}
try:
data = json.loads(self.path.read_text(encoding="utf-8"))
if not isinstance(data, dict) or any(not isinstance(v, dict) or set(v) - FIELDS for v in data.values()):
raise ValueError()
return data
except (ValueError, OSError):
raise RuntimeError("Serverless job journal is unreadable; preserve it for recovery before submitting more work") from None
def all(self):
with self._lock:
return list(self._read().values())
def update(self, record_id, **fields):
if set(fields) - FIELDS:
raise ValueError("Unexpected serverless journal field")
with self._lock:
data = self._read()
data.setdefault(record_id, {"record_id": record_id}).update(fields)
# Keep every unresolved job, plus a bounded terminal history.
completed = sorted((k for k, v in data.items() if v.get("status") in TERMINAL), key=lambda k: data[k].get("updated_at", 0))
for key in completed[:-200]:
data.pop(key)
if self.path.parent.is_symlink():
raise RuntimeError("Unsafe serverless journal directory")
self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
fd, temporary = tempfile.mkstemp(prefix=".jobs-", dir=self.path.parent)
try:
with os.fdopen(fd, "w", encoding="utf-8") as stream:
json.dump(data, stream, allow_nan=False)
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, self.path)
finally:
if os.path.exists(temporary):
os.unlink(temporary)
return dict(data[record_id])
+601
View File
@@ -0,0 +1,601 @@
"""Bounded image transport for generic Runpod-backed ComfyUI jobs."""
import asyncio
import base64
import binascii
import hashlib
import io
import ipaddress
import os
import re
import socket
import stat
import warnings
from pathlib import Path, PurePosixPath
from urllib.parse import urljoin, urlsplit
import aiohttp
import numpy as np
from PIL import Image, UnidentifiedImageError
MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024
MAX_INPUT_TOTAL_BYTES = 10_000_000
MAX_OUTPUT_FILES = 1024
MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024
MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024
MAX_IMAGE_PIXELS = 40_000_000
MAX_OUTPUT_ARRAY_BYTES = 256 * 1024 * 1024
MAX_REDIRECTS = 3
_DOWNLOAD_TIMEOUT_SECONDS = 60
_CHUNK_SIZE = 64 * 1024
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_UNSUPPORTED_IMAGE_EXTENSIONS = {".gif", ".tif", ".tiff", ".avif", ".heic", ".ico", ".svg"}
_AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus", ".aiff", ".aif", ".wma"}
_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
_MEDIA_EXTENSIONS = _IMAGE_EXTENSIONS | _UNSUPPORTED_IMAGE_EXTENSIONS | _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS
_MEDIA_ANNOTATION = re.compile(r"\s+\[([A-Za-z0-9_-]{1,64})\]$")
_PREFIX_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\Z", re.ASCII)
_SAFE_FILENAME_RE = re.compile(r"[A-Za-z0-9_.-]{1,255}\Z", re.ASCII)
class ServerlessMediaError(ValueError):
"""Sanitized input, artifact, or download validation failure."""
def _folder_roots(folder_paths):
roots = {}
for kind, getter_name in (
("input", "get_input_directory"),
("output", "get_output_directory"),
("temp", "get_temp_directory"),
):
getter = getattr(folder_paths, getter_name, None)
if not callable(getter):
raise ServerlessMediaError("ComfyUI media directories are unavailable")
try:
roots[kind] = Path(getter()).resolve(strict=True)
except Exception:
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
return roots
def _inside_roots(path, roots):
for root in roots.values():
try:
path.relative_to(root)
return True
except ValueError:
continue
return False
def _normalize_media_reference(value):
if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value:
raise ServerlessMediaError("Image input reference is invalid")
value = value.replace("\\", "/")
explicit_root = None
annotation = _MEDIA_ANNOTATION.search(value)
if annotation:
marker = annotation.group(1).lower()
if marker not in {"input", "output", "temp"}:
raise ServerlessMediaError("Image input annotation is unsupported")
explicit_root = marker
value = value[: annotation.start()]
if "://" in value:
raise ServerlessMediaError("Remote media references are not supported as image inputs")
if re.match(r"^[A-Za-z]:", value) or value.startswith("//"):
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
path = PurePosixPath(value)
if any(part == ".." for part in path.parts):
raise ServerlessMediaError("Image input path traversal is not allowed")
if not path.parts or value in {".", ""}:
raise ServerlessMediaError("Image input reference is invalid")
if explicit_root:
value = f"{value} [{explicit_root}]"
suffix = Path(path.name).suffix.lower()
return value, suffix
def _resolve_input_path(value, suffix, folder_paths, roots):
if suffix not in _IMAGE_EXTENSIONS:
if suffix in _MEDIA_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
raise ServerlessMediaError("Unsupported serverless image file type")
try:
source = Path(value)
if source.is_absolute():
annotation = _MEDIA_ANNOTATION.search(value)
unannotated = value[:annotation.start()] if annotation else value
resolved = Path(unannotated).resolve(strict=True)
else:
resolver = getattr(folder_paths, "get_annotated_filepath", None)
if not callable(resolver):
raise ServerlessMediaError("ComfyUI media path resolver is unavailable")
resolved = Path(resolver(value)).resolve(strict=True)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if not _inside_roots(resolved, roots):
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
try:
mode = resolved.stat().st_mode
except OSError:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if not stat.S_ISREG(mode):
raise ServerlessMediaError("Image input is not a regular file")
return resolved
def _iter_image_references(workflow):
if not isinstance(workflow, dict):
raise ServerlessMediaError("Serverless workflow must be an API graph")
references = []
for node_id, node in workflow.items():
if not isinstance(node, dict):
continue
class_type = node.get("class_type")
inputs = node.get("inputs", {})
if class_type == "LoadImage" and not isinstance(inputs, dict):
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
if not isinstance(inputs, dict):
continue
if class_type == "LoadImage":
value = inputs.get("image")
if not isinstance(value, str):
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
references.append((node_id, inputs, "image", value, True))
for key in ("image", "file"):
if key == "image" and class_type == "LoadImage":
continue
value = inputs.get(key)
if not isinstance(value, str):
continue
normalized, suffix = _normalize_media_reference(value)
if suffix in _IMAGE_EXTENSIONS:
references.append((node_id, inputs, key, normalized, False))
elif suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
elif suffix in _UNSUPPORTED_IMAGE_EXTENSIONS or (key == "image" and suffix):
raise ServerlessMediaError("Unsupported serverless image file type")
for key in ("audio", "video"):
value = inputs.get(key)
if not isinstance(value, str):
continue
_normalized, suffix = _normalize_media_reference(value)
if suffix in _MEDIA_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
return references
def _read_limited(path, expected_size):
try:
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
with os.fdopen(descriptor, "rb") as handle:
info = os.fstat(handle.fileno())
if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES:
raise ServerlessMediaError("Image input exceeds the permitted size")
if info.st_size > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
data = handle.read(MAX_INPUT_TOTAL_BYTES + 1)
except ServerlessMediaError:
raise
except OSError:
raise ServerlessMediaError("Image input file could not be read safely") from None
if len(data) > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
if len(data) != expected_size:
raise ServerlessMediaError("Image input changed while being read")
return data
async def prepare_images(workflow: dict) -> list[dict]:
"""Pack referenced local images and rewrite only their graph input references.
``workflow`` is expected to be a private API-graph copy owned by the caller.
Returned ``name`` values are flat generated names suitable for ``input.images``.
"""
references = _iter_image_references(workflow)
if not references:
return []
# Import only in ComfyUI, keeping this module importable in the serverless client.
try:
import folder_paths
except Exception:
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
roots = _folder_roots(folder_paths)
resolved_refs = []
unique_paths = {}
total_source_bytes = 0
for node_id, inputs, key, value, _is_load_image in references:
normalized, suffix = _normalize_media_reference(value)
if key == "image" and suffix not in _IMAGE_EXTENSIONS:
if suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
raise ServerlessMediaError("Unsupported serverless image file type")
path = _resolve_input_path(normalized, suffix, folder_paths, roots)
if path not in unique_paths:
try:
size = path.stat().st_size
except OSError:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if size < 0 or size > MAX_INPUT_IMAGE_BYTES:
raise ServerlessMediaError("Image input exceeds the permitted size")
total_source_bytes += size
if total_source_bytes > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size is {total_source_bytes:,} bytes; "
f"the input read guard is {MAX_INPUT_TOTAL_BYTES:,} bytes. Base64 encoding "
"adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
unique_paths[path] = {"size": size, "suffix": suffix}
resolved_refs.append((inputs, key, path))
packed_by_name = {}
path_to_name = {}
for path, metadata in unique_paths.items():
data = _read_limited(path, metadata["size"])
digest = hashlib.sha256(data).hexdigest()
name = f"serverless_{digest}{metadata['suffix']}"
path_to_name[path] = name
packed_by_name.setdefault(name, {"name": name, "image": base64.b64encode(data).decode("ascii")})
for inputs, key, path in resolved_refs:
inputs[key] = path_to_name[path]
return [packed_by_name[name] for name in sorted(packed_by_name)]
def _validate_artifact_url(url):
if not isinstance(url, str) or not url or any(ord(char) < 0x20 or char.isspace() for char in url):
raise ServerlessMediaError("Serverless image artifact URL is invalid")
try:
parsed = urlsplit(url)
port = parsed.port
except ValueError:
raise ServerlessMediaError("Serverless image artifact URL is invalid") from None
if (
parsed.scheme.lower() != "https"
or not parsed.hostname
or parsed.username is not None
or parsed.password is not None
or port not in (None, 443)
or parsed.fragment
):
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
host = parsed.hostname.rstrip(".").lower()
if host in {"localhost", "metadata", "metadata.google.internal"} or host.endswith((".localhost", ".local", ".internal")):
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
try:
address = ipaddress.ip_address(host)
except ValueError:
return url
if not address.is_global:
raise ServerlessMediaError("Serverless image artifact URL must resolve to a public address")
return url
class _PublicResolver(aiohttp.abc.AbstractResolver):
"""Resolve once, reject non-public answers, and pin the checked address."""
async def _getaddrinfo(self, host, port, family):
return await asyncio.get_running_loop().getaddrinfo(
host, port, family=family, type=socket.SOCK_STREAM
)
async def resolve(self, host, port=0, family=socket.AF_INET):
try:
results = await self._getaddrinfo(host, port, family)
except (OSError, asyncio.TimeoutError):
raise OSError("Public artifact host could not be resolved") from None
if not results:
raise OSError("Public artifact host could not be resolved")
resolved = []
seen = set()
for result_family, _socktype, proto, _canonname, sockaddr in results:
try:
address = ipaddress.ip_address(sockaddr[0])
except (ValueError, TypeError, IndexError):
raise OSError("Public artifact host resolved to an invalid address") from None
if not address.is_global:
raise OSError("Public artifact host resolved to a non-public address")
key = (str(address), result_family)
if key in seen:
continue
seen.add(key)
resolved.append({
"hostname": host,
"host": str(address),
"port": port,
"family": result_family,
"proto": proto,
"flags": 0,
})
if not resolved:
raise OSError("Public artifact host could not be resolved")
return resolved
async def close(self):
return None
async def download_public_artifact(url: str, max_bytes: int = MAX_OUTPUT_IMAGE_BYTES) -> bytes:
"""Download a signed public HTTPS artifact without provider credentials."""
if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or not 0 < max_bytes <= MAX_OUTPUT_IMAGE_BYTES:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
current_url = _validate_artifact_url(url)
resolver = _PublicResolver()
connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False, limit=4)
timeout = aiohttp.ClientTimeout(
total=_DOWNLOAD_TIMEOUT_SECONDS,
connect=10,
sock_read=15,
)
try:
async with aiohttp.ClientSession(
connector=connector,
timeout=timeout,
trust_env=False,
cookie_jar=aiohttp.DummyCookieJar(),
) as session:
for redirect_count in range(MAX_REDIRECTS + 1):
_validate_artifact_url(current_url)
async with session.get(current_url, allow_redirects=False) as response:
if response.status in {301, 302, 303, 307, 308}:
location = response.headers.get("Location")
if not location or redirect_count >= MAX_REDIRECTS:
raise ServerlessMediaError("Serverless image artifact redirect is invalid")
current_url = _validate_artifact_url(urljoin(current_url, location))
continue
if response.status != 200:
raise ServerlessMediaError("Serverless image artifact download failed")
payload = bytearray()
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
if not isinstance(chunk, bytes):
raise ServerlessMediaError("Serverless image artifact response is invalid")
if len(payload) + len(chunk) > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
payload.extend(chunk)
if not payload:
raise ServerlessMediaError("Serverless image artifact is empty")
return bytes(payload)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Serverless image artifact download failed") from None
raise ServerlessMediaError("Serverless image artifact download failed")
def _validated_manifest(manifest):
if not isinstance(manifest, list) or len(manifest) > MAX_OUTPUT_FILES:
raise ServerlessMediaError("Serverless image manifest is invalid")
normalized = []
seen_prefixes = set()
prefixes = []
for item in manifest:
if not isinstance(item, dict):
raise ServerlessMediaError("Serverless image manifest is invalid")
if item.get("kind") != "image":
raise ServerlessMediaError("Only image collectors are supported by serverless media transport")
for key in ("collector_id", "job_id", "node_id"):
value = item.get(key)
if not isinstance(value, (str, int)) or isinstance(value, bool) or not str(value):
raise ServerlessMediaError("Serverless image manifest is invalid")
prefix = item.get("prefix")
if not isinstance(prefix, str) or not _PREFIX_RE.fullmatch(prefix) or prefix in {".", ".."}:
raise ServerlessMediaError("Serverless image manifest contains an unsafe filename prefix")
folded = prefix.casefold()
if folded in seen_prefixes:
raise ServerlessMediaError("Serverless image manifest prefixes must be unique")
seen_prefixes.add(folded)
prefixes.append(prefix)
normalized.append({
"collector_id": str(item["collector_id"]),
"job_id": str(item["job_id"]),
"node_id": str(item["node_id"]),
"prefix": prefix,
})
for first in prefixes:
for second in prefixes:
if first != second and second.casefold().startswith(first.casefold() + "_"):
raise ServerlessMediaError("Serverless image manifest prefixes are ambiguous")
return normalized
def _artifact_filename_matches(filename, prefixes):
if not isinstance(filename, str):
return None
basename = filename.replace("\\", "/").rsplit("/", 1)[-1]
matches = [item for item in prefixes if basename.startswith(item["prefix"] + "_")]
if len(matches) > 1:
raise ServerlessMediaError("Serverless image artifact matches multiple collectors")
return matches[0] if matches else None
def _canonical_artifact_identity(entry, filename):
artifact_path = entry.get("artifact_path")
artifact_type = entry.get("artifact_type", "output")
if artifact_path is not None:
if not isinstance(artifact_path, str) or not artifact_path or "\x00" in artifact_path:
raise ServerlessMediaError("Serverless image artifact identity is invalid")
normalized_path = os.path.normcase(os.path.normpath(artifact_path.replace("\\", "/")))
else:
subfolder = entry.get("subfolder", "")
if not isinstance(subfolder, str) or "\x00" in subfolder:
raise ServerlessMediaError("Serverless image artifact identity is invalid")
normalized_path = os.path.normcase(os.path.normpath(f"{subfolder}/{filename}"))
return str(artifact_type), normalized_path
def _decode_base64_data(value, max_bytes=MAX_OUTPUT_IMAGE_BYTES):
if not isinstance(value, str) or not value:
raise ServerlessMediaError("Serverless image artifact data is invalid")
max_bytes = min(max_bytes, MAX_OUTPUT_IMAGE_BYTES)
max_encoded = ((max_bytes + 2) // 3) * 4
if len(value) > max_encoded:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
padding = len(value) - len(value.rstrip("="))
estimated_size = (len(value) // 4) * 3 - padding
if estimated_size > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
try:
payload = base64.b64decode(value, validate=True)
except (binascii.Error, ValueError):
raise ServerlessMediaError("Serverless image artifact data is invalid") from None
if not payload:
raise ServerlessMediaError("Serverless image artifact is empty")
if len(payload) > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
return payload
def _decode_image_array(payload, total_array_bytes):
try:
with warnings.catch_warnings():
warnings.simplefilter("error", Image.DecompressionBombWarning)
image = Image.open(io.BytesIO(payload))
width, height = image.size
if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
raise ServerlessMediaError("Serverless image dimensions exceed the permitted limit")
has_alpha = image.mode in {"RGBA", "LA"} or "transparency" in image.info
mode = "RGBA" if has_alpha else "RGB"
output_bytes = width * height * (4 if has_alpha else 3) * 4
if total_array_bytes + output_bytes > MAX_OUTPUT_ARRAY_BYTES:
raise ServerlessMediaError("Decoded serverless images exceed the permitted memory limit")
converted = image.convert(mode)
converted.load()
array = np.array(converted, dtype=np.float32, copy=True)
np.multiply(array, 1.0 / 255.0, out=array)
return array, output_bytes
except ServerlessMediaError:
raise
except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
except Exception:
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
def _image_output(output):
"""Unwrap handler output and reject explicit execution errors."""
if not isinstance(output, dict):
raise ServerlessMediaError("Serverless job output is invalid")
nested = output.get("output")
data = nested if isinstance(nested, dict) else output
for candidate in (output, data):
if candidate.get("error") or candidate.get("errors"):
raise ServerlessMediaError("Serverless job reported an execution error")
return data
async def _select_image_payloads(data, expected, download):
"""Match artifact identities and fetch bounded bytes, without image decoding."""
artifacts = data.get("images", [])
if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES:
raise ServerlessMediaError("Serverless job returned too many image files")
selected = {item["prefix"]: [] for item in expected}
identities = set()
total_payload_bytes = 0
for entry in artifacts:
if not isinstance(entry, dict):
continue
filename = entry.get("filename")
manifest_entry = _artifact_filename_matches(filename, expected)
if manifest_entry is None:
continue
normalized_filename = filename.replace("\\", "/")
if "/" in normalized_filename or not _SAFE_FILENAME_RE.fullmatch(filename):
raise ServerlessMediaError("Serverless image artifact filename is invalid")
if Path(filename).suffix.lower() not in _IMAGE_EXTENSIONS:
raise ServerlessMediaError("Serverless image artifact has an unsupported file type")
if "artifact_type" in entry and entry["artifact_type"] != "output":
raise ServerlessMediaError("Expected serverless image artifact was not an output image")
identity = _canonical_artifact_identity(entry, filename)
if identity in identities:
raise ServerlessMediaError("Duplicate serverless image artifact identity")
identities.add(identity)
match = re.match(rf"^{re.escape(manifest_entry['prefix'])}_(\d+)(?:_|\.)", filename)
if match is None:
raise ServerlessMediaError("Expected serverless image filename counter is invalid")
artifact_type = entry.get("type")
remaining_bytes = MAX_OUTPUT_TOTAL_BYTES - total_payload_bytes
if remaining_bytes <= 0:
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
if artifact_type == "base64":
payload = _decode_base64_data(entry.get("data"), remaining_bytes)
elif artifact_type in {"gcs_url", "s3_url", "url"}:
url = entry.get("data") or entry.get("url")
url = _validate_artifact_url(url)
try:
if download is None:
payload = await download_public_artifact(
url, max_bytes=min(MAX_OUTPUT_IMAGE_BYTES, remaining_bytes)
)
else:
payload = await download(url)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Serverless image artifact download failed") from None
if not isinstance(payload, bytes) or not payload:
raise ServerlessMediaError("Serverless image artifact response is invalid")
if len(payload) > MAX_OUTPUT_IMAGE_BYTES:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
else:
raise ServerlessMediaError("Serverless image artifact transport type is unsupported")
total_payload_bytes += len(payload)
if total_payload_bytes > MAX_OUTPUT_TOTAL_BYTES:
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
selected[manifest_entry["prefix"]].append((int(match.group(1)), filename, identity, payload))
if any(not selected[item["prefix"]] for item in expected):
raise ServerlessMediaError("Serverless job is missing an expected collector image")
return selected
def _decode_image_groups(selected, expected):
"""Decode in stable counter order, bounding array memory across collectors."""
result = {}
total_array_bytes = 0
for manifest_entry in expected:
prefix = manifest_entry["prefix"]
rows = sorted(selected[prefix], key=lambda row: (row[0], row[1], row[2]))
arrays = []
shape = None
for _counter, _filename, _identity, payload in rows:
array, added_bytes = _decode_image_array(payload, total_array_bytes)
total_array_bytes += added_bytes
if shape is None:
shape = array.shape
elif array.shape != shape:
raise ServerlessMediaError("Images from one serverless collector have inconsistent dimensions")
arrays.append(array)
result.setdefault(manifest_entry["job_id"], []).extend(arrays)
return result
async def collect_images(output: dict, manifest: list, download=None) -> dict[str, list[np.ndarray]]:
"""Validate and decode expected image artifacts, grouped by serverless job ID."""
data = _image_output(output)
expected = _validated_manifest(manifest)
if not expected:
return {}
selected = await _select_image_payloads(data, expected, download)
return _decode_image_groups(selected, expected)
@@ -0,0 +1,199 @@
"""Pure compilation of prepared typed-media collectors for serverless workers.
The legacy image compiler remains separate. This compiler emits one MediaExportV1
execution target per connected socket, retaining only its upstream computation.
It never imports ComfyUI or inspects/executes runtime node classes.
"""
from __future__ import annotations
import copy
import hashlib
import json
from graphlib import CycleError, TopologicalSorter
from ...serverless_bridge.contract import SOCKET_KINDS, identifier, validate_bindings
from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
from .serverless_prompt import (
_PRIMITIVE_NODES,
_REMOVED_SCALAR_NODES,
_is_link,
_is_unsupported_distributed,
_make_scalar_resolver,
_prune_graph,
)
_COLLECTOR = "DistributedMediaCollector"
def _validate_identifier(value, context):
# Validate the original token before any lookup. In particular, neither
# integer node IDs nor whitespace in prepared identities is normalized.
try:
identifier(value)
except ValueError:
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
def _validate_prompt(prepared_prompt):
"""Validate the entire source graph, including branches later discarded."""
if not isinstance(prepared_prompt, dict) or not prepared_prompt:
raise ValueError("prepared prompt must be a nonempty API graph mapping")
collectors = []
for node_id, node in prepared_prompt.items():
_validate_identifier(node_id, "prompt node ID")
if not isinstance(node, dict):
raise ValueError(f"prompt node {node_id!r} must be an object")
class_type = node.get("class_type")
if not isinstance(class_type, str) or not class_type.strip():
raise ValueError(f"prompt node {node_id!r} must have a nonempty class_type")
if class_type == _COLLECTOR:
collectors.append(node_id)
elif class_type == "DistributedCollector" or _is_unsupported_distributed(class_type):
raise ValueError(f"unsupported distributed or batch node: {class_type}")
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
if any(not isinstance(name, str) or not name for name in inputs):
raise ValueError(f"prompt node {node_id!r} input names must be nonempty strings")
if not collectors:
raise ValueError("prepared prompt contains no DistributedMediaCollector nodes")
dependencies = {}
for node_id, node in prepared_prompt.items():
parents = set()
for input_name, value in node.get("inputs", {}).items():
# Ordinary inputs may contain literal lists. As in the legacy
# compiler, exactly two elements denotes an API link, not media.
if not _is_link(value):
continue
source_id, slot = value
context = f"link at node {node_id!r} input {input_name!r}"
_validate_identifier(source_id, context + " source ID")
if source_id not in prepared_prompt:
raise ValueError(f"invalid {context}: source {source_id!r} does not exist")
if type(slot) is not int or slot < 0:
raise ValueError(f"invalid {context}: output slot must be a nonnegative integer")
source_class = prepared_prompt[source_id]["class_type"]
if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
if slot != 0:
raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
elif source_class == _COLLECTOR and slot >= len(SOCKET_KINDS):
raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
parents.add(source_id)
dependencies[node_id] = parents
for collector_id in collectors:
inputs = prepared_prompt[collector_id].get("inputs", {})
validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
inputs.get("transfer_quality", DEFAULT_QUALITY))
for field in ("multi_job_id", "worker_id"):
_validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}")
connected = [socket for socket in SOCKET_KINDS if socket in inputs]
if not connected:
raise ValueError(f"collector {collector_id!r} requires at least one connected media socket")
for socket in connected:
value = inputs[socket]
if not _is_link(value):
raise ValueError(f"collector {collector_id!r} {socket} must be linked; use PrimitiveString for literal text")
if prepared_prompt[value[0]]["class_type"] in _REMOVED_SCALAR_NODES:
# Export sockets must remain actual links after scalar folding.
# A PrimitiveString wrapper supports distributed text values.
raise ValueError(f"collector {collector_id!r} {socket} cannot link directly to a distributed scalar; use a typed node or PrimitiveString")
# Iterative topological sorting checks ordinary and unused cycles too, not
# only the primitive chains that the shared scalar resolver traverses.
try:
ordered = list(TopologicalSorter(dependencies).static_order())
except CycleError:
raise ValueError("cycle in prepared media prompt") from None
collector_set = set(collectors)
has_collector_ancestor = {}
for node_id in ordered:
has_collector_ancestor[node_id] = any(
parent in collector_set or has_collector_ancestor[parent]
for parent in dependencies[node_id]
)
if node_id in collector_set and has_collector_ancestor[node_id]:
raise ValueError("collector-to-collector dependence is unsupported")
return sorted(collectors)
def _build_exports(source_nodes, collector_ids, request_id):
bindings = []
exports = {}
reserved = set(source_nodes)
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
for socket, kind in SOCKET_KINDS.items():
if socket not in inputs:
continue
# Hashing avoids separator ambiguity and keeps long valid source
# identities within the bridge's identifier length bound. Identity
# is stable across request IDs and dictionary insertion orders.
identity = json.dumps([collector_id, socket], separators=(",", ":"))
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
base_id = f"serverless:media:{digest}"
node_id = base_id
suffix = 0
while node_id in reserved:
suffix += 1
node_id = f"{base_id}:{suffix}"
reserved.add(node_id)
binding = {
"binding_id": f"b:{digest}",
"collector_id": collector_id,
"job_id": inputs["multi_job_id"],
"participant_id": inputs["worker_id"],
"node_id": node_id,
"socket": socket,
"kind": kind,
}
bindings.append(binding)
exports[node_id] = {
"class_type": "MediaExportV1",
"inputs": {
"request_id": request_id,
"binding_json": json.dumps(binding),
socket: copy.deepcopy(inputs[socket]),
},
}
if kind == "IMAGE":
fmt, quality = validate_options(inputs.get("transfer_format", DEFAULT_FORMAT),
inputs.get("transfer_quality", DEFAULT_QUALITY))
if fmt != DEFAULT_FORMAT:
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
validate_bindings(bindings)
return bindings, exports
def compile_media_prompt(prepared_prompt: dict, request_id: str) -> tuple[dict, list[dict]]:
"""Return a worker graph and canonical typed-media export bindings.
Each DistributedMediaCollector must have prepared literal ``multi_job_id``
and ``worker_id`` identities and at least one linked media input. Present
media inputs (including texts) must be API links; literal text belongs in a
PrimitiveString node. Nonmedia collector controls are not exported.
Only MediaExportV1 targets and their ancestors survive. Required ordinary
output/utility nodes are preserved regardless of their runtime OUTPUT_NODE
flag; callers must dispatch exactly the returned binding ``node_id`` values
as partial_execution_targets. Distributed scalar dependencies use the same
resolver as the legacy compiler. Direct scalar-to-media links need a typed
wrapper so that each export input remains a link.
Validation and compilation never modify the caller's source graph. This
structural compiler does not infer socket types or unknown nodes' output
counts; ComfyUI performs runtime type validation at the endpoint.
"""
_validate_identifier(request_id, "request_id")
collector_ids = _validate_prompt(prepared_prompt)
bindings, exports = _build_exports(prepared_prompt, collector_ids, request_id)
compiled = _prune_graph(
prepared_prompt, bindings, exports, _make_scalar_resolver(prepared_prompt)
)
return {node_id: compiled[node_id] for node_id in sorted(compiled)}, bindings
+420
View File
@@ -0,0 +1,420 @@
"""Compile prepared distributed prompts into worker-side ComfyUI image jobs.
This module is deliberately pure: it only copies and rewrites an API prompt
mapping and never imports or executes ComfyUI nodes.
"""
from __future__ import annotations
import copy
import hashlib
import json
import re
_SUPPORTED_DISTRIBUTED = {
"DistributedCollector",
"DistributedSeed",
"DistributedValue",
"DistributedModelName",
}
_REMOVED_SCALAR_NODES = {
"DistributedSeed",
"DistributedValue",
"DistributedModelName",
}
_PRIMITIVE_NODES = {
"PrimitiveBoolean": False,
"PrimitiveFloat": 0.0,
"PrimitiveInt": 0,
"PrimitiveString": "",
}
_REMOVED_NODE_OUTPUTS = {
"DistributedCollector": {0, 1},
"DistributedSeed": {0},
"DistributedValue": {0},
"DistributedModelName": {0},
}
_SEED_DEFAULT = 1125899906842
def _is_link(value):
return isinstance(value, list) and len(value) == 2
def _is_unsupported_distributed(class_type):
if not isinstance(class_type, str) or class_type in _SUPPORTED_DISTRIBUTED:
return False
lowered = class_type.lower()
return "distributed" in lowered or "batchdivider" in lowered
def _safe_component(value):
cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_-")
return cleaned[:24] or "empty"
def _filename_prefix(job_id, worker_id, collector_id, kind):
identity = json.dumps(
[job_id, worker_id, collector_id, kind],
ensure_ascii=False,
separators=(",", ":"),
)
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
parts = (
"serverless",
_safe_component(job_id),
_safe_component(worker_id),
_safe_component(collector_id),
kind,
digest,
)
return "_".join(parts)
def _worker_index(worker_id):
"""Mirror nodes.utilities worker-index parsing; None means fallback."""
try:
if worker_id.startswith("worker_"):
return int(worker_id.split("_")[1])
return int(worker_id)
except (ValueError, IndexError):
return None
def _coerce(value, value_type):
if value_type == "INT":
return int(float(value))
if value_type == "FLOAT":
return float(value)
return value
def _coerce_safe(value, value_type):
try:
return _coerce(value, value_type)
except (TypeError, ValueError):
return value
def _stringify_model_name(value):
if isinstance(value, str):
return value
if isinstance(value, (int, float, bool)):
return str(value)
try:
return json.dumps(value, indent=4)
except (TypeError, ValueError):
return str(value)
def _model_name_result(text):
if isinstance(text, list):
values = [_stringify_model_name(value) for value in text]
return values[0] if len(values) == 1 else values
return _stringify_model_name(text)
def _validate_prepared_prompt(prepared_prompt):
"""Validate graph shape, links, collectors, and supported topology."""
if not isinstance(prepared_prompt, dict):
raise ValueError("prepared prompt must be an API graph mapping")
source_nodes = {}
for raw_id, node in prepared_prompt.items():
node_id = str(raw_id)
if node_id in source_nodes:
raise ValueError("prompt has duplicate node IDs after string normalization")
if not isinstance(node, dict):
raise ValueError(f"prompt node {node_id!r} must be an object")
class_type = node.get("class_type")
if not isinstance(class_type, str):
raise ValueError(f"prompt node {node_id!r} is missing a class_type")
if _is_unsupported_distributed(class_type):
raise ValueError(f"unsupported distributed or batch node: {class_type}")
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
source_nodes[node_id] = node
if not source_nodes:
raise ValueError("prepared prompt contains no nodes")
# Validate every API link before any graph rewriting. Slots on the nodes
# removed by this compiler have a known schema and are checked exactly.
for node_id, node in source_nodes.items():
for input_name, value in node.get("inputs", {}).items():
if not _is_link(value):
continue
upstream_id, output_slot = value
upstream_id = str(upstream_id)
if upstream_id not in source_nodes:
raise ValueError(
f"invalid link at node {node_id!r} input {input_name!r}: "
f"source {upstream_id!r} does not exist"
)
if isinstance(output_slot, bool) or not isinstance(output_slot, int) or output_slot < 0:
raise ValueError(
f"invalid link at node {node_id!r} input {input_name!r}: "
"output slot must be a nonnegative integer"
)
upstream_class = source_nodes[upstream_id]["class_type"]
valid_slots = _REMOVED_NODE_OUTPUTS.get(upstream_class)
if valid_slots is not None and output_slot not in valid_slots:
raise ValueError(
f"invalid output slot {output_slot} for removed node "
f"{upstream_class} {upstream_id!r}"
)
collector_ids = sorted(
(node_id for node_id, node in source_nodes.items() if node["class_type"] == "DistributedCollector"),
key=str,
)
if not collector_ids:
raise ValueError("prepared prompt contains no DistributedCollector nodes")
# Only image collectors are supported by the current generic handler.
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
if "audio" in inputs:
raise ValueError("collector audio input is unsupported by this image-only compiler")
if not _is_link(inputs.get("images")):
raise ValueError(f"collector {collector_id!r} is missing a linked image input")
source_id = str(inputs["images"][0])
source_class = source_nodes[source_id]["class_type"]
if source_class in _REMOVED_SCALAR_NODES:
raise ValueError(f"collector {collector_id!r} image input must link to an image node")
# A collector cannot be used, directly or through processing, as another
# collector's input: this first compiler handles disconnected jobs only.
for collector_id in collector_ids:
pending = [str(source_nodes[collector_id]["inputs"]["images"][0])]
seen = set()
while pending:
node_id = pending.pop()
if node_id in seen:
continue
seen.add(node_id)
if source_nodes[node_id]["class_type"] == "DistributedCollector":
raise ValueError("collector-to-collector dependence is unsupported")
for value in source_nodes[node_id].get("inputs", {}).values():
if _is_link(value):
pending.append(str(value[0]))
return source_nodes, collector_ids
def _make_scalar_resolver(source_nodes):
"""Build pure resolvers for distributed scalar nodes and stock primitives."""
# Resolve only stock primitive nodes. A visited set makes malformed
# primitive-to-primitive cycles fail clearly instead of recursing forever.
def resolve_primitive(node_id, visiting=None):
node_id = str(node_id)
node = source_nodes[node_id]
class_type = node["class_type"]
if class_type not in _PRIMITIVE_NODES:
raise ValueError("linked value must be literal or supported PrimitiveInt/Float/String/Boolean")
if visiting is None:
visiting = set()
if node_id in visiting:
raise ValueError(f"cycle while resolving primitive value at node {node_id!r}")
visiting.add(node_id)
value = node.get("inputs", {}).get("value", _PRIMITIVE_NODES[class_type])
if _is_link(value):
parent_id, output_slot = str(value[0]), value[1]
if output_slot != 0:
raise ValueError(f"invalid output slot {output_slot} for primitive node {parent_id!r}")
if source_nodes[parent_id]["class_type"] not in _PRIMITIVE_NODES:
raise ValueError("linked primitive value must come from a stock Primitive node")
result = resolve_primitive(parent_id, visiting)
else:
result = copy.deepcopy(value)
visiting.remove(node_id)
return result
def resolve_literal_or_primitive(value, context):
if _is_link(value):
source_id, output_slot = str(value[0]), value[1]
if output_slot != 0 or source_nodes[source_id]["class_type"] not in _PRIMITIVE_NODES:
raise ValueError(f"{context} must be a literal or supported Primitive link")
try:
return resolve_primitive(source_id)
except ValueError as exc:
if "cycle" in str(exc):
raise
raise ValueError(f"{context} must be a literal or supported Primitive link") from None
return copy.deepcopy(value)
def resolve_hidden(node_id, name, default):
value = source_nodes[node_id].get("inputs", {}).get(name, default)
if _is_link(value):
raise ValueError(f"{name} on distributed node {node_id!r} must be a literal")
return copy.deepcopy(value)
resolved_distributed = {}
def resolve_distributed(node_id, output_slot):
node_id = str(node_id)
class_type = source_nodes[node_id]["class_type"]
if output_slot != 0:
raise ValueError(f"invalid output slot {output_slot} for {class_type} {node_id!r}")
cache_key = (node_id, output_slot)
if cache_key in resolved_distributed:
return copy.deepcopy(resolved_distributed[cache_key])
inputs = source_nodes[node_id].get("inputs", {})
if class_type == "DistributedSeed":
seed = resolve_literal_or_primitive(inputs.get("seed", _SEED_DEFAULT), f"seed on {node_id!r}")
is_worker = resolve_hidden(node_id, "is_worker", False)
worker_id = resolve_hidden(node_id, "worker_id", "")
if is_worker:
worker_index = _worker_index(worker_id)
if worker_index is not None:
seed = seed + worker_index + 1
result = seed
elif class_type == "DistributedValue":
default_value = resolve_literal_or_primitive(
inputs.get("default_value", ""), f"default_value on {node_id!r}"
)
worker_values = resolve_literal_or_primitive(
inputs.get("worker_values", "{}"), f"worker_values on {node_id!r}"
)
is_worker = resolve_hidden(node_id, "is_worker", False)
worker_id = resolve_hidden(node_id, "worker_id", "")
try:
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
except json.JSONDecodeError:
values = {}
if not isinstance(values, dict):
values = {}
value_type = values.get("_type", "STRING")
fallback = _coerce_safe(default_value, value_type)
result = fallback
if is_worker:
worker_index = _worker_index(worker_id)
if worker_index is not None:
raw_value = values.get(str(worker_index + 1), "")
if raw_value:
try:
result = _coerce(raw_value, value_type)
except (ValueError, IndexError):
result = fallback
elif class_type == "DistributedModelName":
text = resolve_literal_or_primitive(
inputs.get("text", ""), f"DistributedModelName text on {node_id!r}"
)
result = _model_name_result(text)
else:
raise ValueError(f"unsupported distributed node: {class_type}")
resolved_distributed[cache_key] = copy.deepcopy(result)
return result
return resolve_distributed
def _build_savers(source_nodes, collector_ids):
"""Construct stable SaveImage nodes and their collector metadata."""
# The prepared collector values define the artifact identity. The fallback
# for absent hidden values is stable and retains the collector's job scope.
collector_specs = []
max_numeric_id = 0
for node_id in source_nodes:
try:
max_numeric_id = max(max_numeric_id, int(node_id))
except ValueError:
pass
next_numeric_id = max_numeric_id + 1
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
raw_job_id = inputs.get("multi_job_id", collector_id)
if _is_link(raw_job_id):
raise ValueError(f"collector {collector_id!r} multi_job_id must be a literal")
job_id = str(collector_id if raw_job_id in (None, "") else raw_job_id)
raw_worker_id = inputs.get("worker_id", "")
if _is_link(raw_worker_id):
raise ValueError(f"collector {collector_id!r} worker_id must be a literal")
worker_id = "" if raw_worker_id is None else str(raw_worker_id)
saver_id = str(next_numeric_id)
next_numeric_id += 1
prefix = _filename_prefix(job_id, worker_id, collector_id, "image")
collector_specs.append(
{
"collector_id": collector_id,
"job_id": job_id,
"node_id": saver_id,
"kind": "image",
"prefix": prefix,
"source_link": copy.deepcopy(inputs["images"]),
}
)
saver_nodes = {}
for spec in collector_specs:
saver_nodes[spec["node_id"]] = {
"class_type": "SaveImage",
"inputs": {
"images": spec["source_link"],
"filename_prefix": spec["prefix"],
},
"_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"},
}
return collector_specs, saver_nodes
def _prune_graph(source_nodes, collector_specs, saver_nodes, resolve_distributed):
"""Keep saver ancestors, replacing distributed scalar links with literals."""
all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()}
all_nodes.update(copy.deepcopy(saver_nodes))
needed = set()
pending = [spec["node_id"] for spec in collector_specs]
while pending:
node_id = pending.pop()
if node_id in needed:
continue
needed.add(node_id)
node = all_nodes[node_id]
if node["class_type"] == "DistributedCollector":
raise ValueError("collector-to-collector dependence is unsupported")
for input_name, value in list(node.get("inputs", {}).items()):
if not _is_link(value):
continue
upstream_id, output_slot = str(value[0]), value[1]
upstream_class = source_nodes[upstream_id]["class_type"]
if upstream_class in _REMOVED_SCALAR_NODES:
node["inputs"][input_name] = resolve_distributed(upstream_id, output_slot)
else:
pending.append(upstream_id)
compiled = {
node_id: copy.deepcopy(node)
for node_id, node in all_nodes.items()
if node_id in needed
}
if any("distributed" in node["class_type"].lower() for node in compiled.values()):
raise ValueError("compiled prompt still contains a distributed dependency")
return compiled
def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]:
"""Compile a prepared worker prompt to an image-only serverless graph.
Every DistributedCollector's original ``images`` input is connected to a
new SaveImage node. Only those savers and their upstream dependencies are
retained. DistributedSeed/Value/ModelName links are resolved as scalar
literals; unsupported node execution is never attempted.
"""
source_nodes, collector_ids = _validate_prepared_prompt(prepared_prompt)
resolve_distributed = _make_scalar_resolver(source_nodes)
collector_specs, saver_nodes = _build_savers(source_nodes, collector_ids)
compiled = _prune_graph(
source_nodes,
collector_specs,
saver_nodes,
resolve_distributed,
)
manifest = [
{key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")}
for spec in collector_specs
]
return compiled, manifest
+413
View File
@@ -0,0 +1,413 @@
"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
import asyncio
import time
import uuid
from pathlib import Path
from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
from .serverless_ledger import JobLedger
from .serverless_state import (
PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
occupies_capacity, terminal_provider_status,
)
# Fixed per-worker cap on this master's outstanding Runpod jobs; equals the
# Images per run maximum (MAX_JOBS_PER_RUN in utils/serverless_config.py), so
# every image of a run is submitted at once. Deliberately not configurable or
# derived from jobs_per_run: unresolved jobs from earlier runs (ambiguous
# submissions, unconfirmed cancels) still count against it, and a constant
# never has to change while such jobs exist.
PER_WORKER_JOB_LIMIT = 16
async def decode_results(output, manifest):
from .serverless_media import collect_images
import torch
images = await collect_images(output, manifest)
return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
def state_directory():
import folder_paths
return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
def get_secret_store():
from ...utils.serverless_config import SecretStore
return SecretStore(state_directory() / "credentials.json")
def get_manager():
import server
from ...utils.network import get_client_session
ps = server.PromptServer.instance
manager = getattr(ps, "distributed_serverless", None)
if manager is None:
manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
ps.distributed_serverless = manager
return manager
class ServerlessManager:
def __init__(self, prompt_server, secrets, ledger_path, session_factory):
self.server = prompt_server
self.secrets = secrets
self.ledger = JobLedger(ledger_path)
self.session_factory = session_factory
self.tasks = {}
self.runs = {}
self.pool_limits = {}
self.health_cache = {}
self.health_locks = {}
self.poll_seconds = 2.0
self.capacity_changed = asyncio.Condition()
self.recovery_lock = asyncio.Lock()
self._recover_local_records()
def _recover_local_records(self):
# WAITING is committed before admission; it cannot have reached Runpod.
# SUBMITTING is an intent record, so a lost ID must remain reserved.
for record in self.ledger.all():
if record.get("status") == "WAITING" and not record.get("job_id"):
self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time())
elif record.get("status") == "SUBMITTING" and not record.get("job_id"):
self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time())
def _task_is_active(self, record_id):
task = self.tasks.get(record_id)
return task is not None and not task.done()
def pool_is_active(self, worker_id):
"""Protect both live local work and unresolved pre-restart provider work."""
for rid, run in self.runs.items():
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
if str(pool_id) == str(worker_id) and self._task_is_active(rid):
return True
return any(belongs_to_pool(record, worker_id) and is_unresolved(record)
for record in self.ledger.all())
async def _reserve_provider_capacity(self, run):
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
async with self.capacity_changed:
while True:
self._check_master(run)
outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record)
for record in self.ledger.all())
if outstanding < PER_WORKER_JOB_LIMIT:
# Claim atomically before any await that could submit a billed job.
self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time())
return
await self.capacity_changed.wait()
async def _notify_capacity(self):
async with self.capacity_changed:
self.capacity_changed.notify_all()
async def _client(self, worker):
pool_id = worker.get("pool_id", worker["id"])
key = self.secrets.get(pool_id)
return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
async def health(self, worker, *, fresh=False):
pool_id = worker.get("pool_id", worker["id"])
cache_key = (pool_id, worker["endpoint_id"])
lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
async with lock:
cached = self.health_cache.get(cache_key)
if cached and not fresh and time.monotonic() - cached[0] < 10:
return cached[1]
client = await self._client(worker)
data = await client.health()
jobs = data.get("jobs", {})
if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
raise RunpodError("Runpod returned an invalid endpoint health record")
try:
pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
except (ValueError, TypeError):
raise RunpodError("Runpod returned invalid queue counts") from None
result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
self.health_cache[cache_key] = (time.monotonic(), result)
return result
def _semaphore(self, worker):
# The limit is a constant, so it never has to be swapped while this
# worker still has unresolved jobs.
pool_id = worker.get("pool_id", worker["id"])
semaphore = self.pool_limits.get(pool_id)
if semaphore is None:
semaphore = asyncio.Semaphore(PER_WORKER_JOB_LIMIT)
self.pool_limits[pool_id] = semaphore
return semaphore
def ensure_capacity(self, count):
if sum(not task.done() for task in self.tasks.values()) + count > 64:
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
async def start(self, execution_id, prompt_id, prepared_jobs):
self.ensure_capacity(len(prepared_jobs))
from .serverless_typed_results import validate_prepared_job
# Validate the entire batch before creating even the first local slot.
for prepared in prepared_jobs:
validate_prepared_job(prepared)
new_runs = []
# Journal every local slot before starting any billed submission.
for prepared in prepared_jobs:
worker = prepared["worker"]
semaphore = self._semaphore(worker)
record_id = uuid.uuid4().hex
now = time.time()
self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
self.runs[record_id] = run
new_runs.append((record_id, run, semaphore))
for record_id, run, semaphore in new_runs:
self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
# Bound in-memory completed tracking as well as the on-disk journal.
done = [rid for rid, task in self.tasks.items() if task.done()]
for rid in done[:-200]:
self.tasks.pop(rid, None)
self.runs.pop(rid, None)
# Additive ownership handle for item schedulers; legacy callers ignore it.
return [record_id for record_id, _run, _semaphore in new_runs]
def is_pending(self, collector_job_id, participant_id):
return any(not task.done() and run["worker"]["id"] == participant_id and
any(m["job_id"] == collector_job_id for m in run["manifest"])
for rid, task in self.tasks.items() for run in [self.runs[rid]])
def owns(self, collector_job_id, participant_id):
return any(run["worker"]["id"] == participant_id and
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
async def _deliver(self, run, images=None, error=None, *, media=None):
if media is not None and not error:
# Validate all destinations before exposing any partial success.
for binding in run["manifest"]:
if binding["participant_id"] != run["worker"]["id"]:
raise ValueError("Serverless media participant identity mismatch")
values = media[binding["job_id"]][binding["socket"]]
if not isinstance(values, list) or not values:
raise ValueError("Missing serverless media values")
async with self.server.distributed_jobs_lock:
delivered = set()
for binding in run["manifest"]:
job_id = binding["job_id"]
if error or media is not None:
if job_id in delivered:
continue
delivered.add(job_id)
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
if queue is None:
continue # A cancelled/finished collector must never be resurrected.
if error:
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
elif media is not None:
queue.put_nowait({"worker_id": run["worker"]["id"], "media": media[job_id], "is_last": True})
else:
batch = images[binding["job_id"]]
for index, tensor in enumerate(batch):
queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
"tensor": tensor, "is_last": index == len(batch) - 1})
async def _cancel_provider(self, run):
if run["provider_job_id"] and run["client"]:
try:
result = await run["client"].cancel(run["provider_job_id"])
return terminal_provider_status(result)
except Exception:
return None
return None
def _check_master(self, run):
queue = getattr(self.server, "prompt_queue", None)
if queue is None:
return
history = queue.get_history(run["prompt_id"])
record = history.get(run["prompt_id"], {})
if record.get("status", {}).get("status_str") == "error":
raise RunpodError("Master workflow failed; serverless work stopped")
if not record and hasattr(queue, "get_current_queue"):
running, pending = queue.get_current_queue()
if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
async def _run(self, run, semaphore):
rid = run["record_id"]
submitted = False
provider_terminal = False
try:
from .serverless_typed_results import validate_prepared_job, collect_typed_results
typed = validate_prepared_job(run)
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
async with semaphore:
await self._reserve_provider_capacity(run)
client = await self._client(run["worker"])
run["client"] = client
submitted = True
if typed:
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"], media=run["media"])
else:
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
run["provider_job_id"] = job_id
# Persist immediately: all subsequent actions use this exact ID.
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
while True:
self._check_master(run)
response = await client.status(job_id)
status = response.get("status")
if not isinstance(status, str) or status not in PROVIDER_STATES:
raise RunpodError("Runpod returned an unknown job state")
provider_terminal = status in PROVIDER_TERMINAL
if status == "COMPLETED":
output = response.get("output")
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
if typed:
media = await collect_typed_results(output, run["media"], run["manifest"])
await self._deliver(run, media=media)
else:
images = await decode_results(output, run["manifest"])
await self._deliver(run, images=images)
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
return
self.ledger.update(rid, status=status, updated_at=time.time())
if provider_terminal:
raise RunpodError(f"Serverless job ended with {status}")
await asyncio.sleep(self.poll_seconds)
except AmbiguousSubmission as exc:
self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time())
await self._deliver(run, error=str(exc))
except (asyncio.TimeoutError, asyncio.CancelledError) as exc:
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
if submitted and not run["provider_job_id"]:
status = "UNKNOWN"
error = "Submission interrupted before its job ID was received; check Runpod before retrying"
elif isinstance(exc, asyncio.TimeoutError):
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
error = "Serverless job exceeded its queue/execution deadline"
else:
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED")
error = "Serverless work stopped" if terminal or not submitted else "Serverless cancellation requested; provider confirmation is pending"
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
await self._deliver(run, error=error)
except Exception as exc:
# If polling/decoding failed, stop only our own provider job, never the pool.
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract"
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED")
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
await self._deliver(run, error=error)
finally:
# Completed tracking retains IDs/bindings only, not potentially large
# base64 inputs, workflow metadata, or a client holding credentials.
run.pop("workflow", None)
run.pop("images", None)
run.pop("media", None)
run["client"] = None
# A released local semaphore alone cannot release a durable reservation.
await self._notify_capacity()
async def wait_for_execution(self, execution_id):
tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id]
await asyncio.gather(*tasks, return_exceptions=False)
def _detached_records(self, worker_id=None, execution_id=None):
return [record for record in self.ledger.all()
if is_unresolved(record) and not self._task_is_active(record["record_id"])
and (worker_id is None or belongs_to_pool(record, worker_id))
and (execution_id is None or record.get("execution_id") == execution_id)]
async def _recover_provider_record(self, record, *, cancel=False):
"""Reconcile an exact saved ID, never submit or resume a workflow."""
if not record.get("job_id"):
return False
status = None
try:
client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]})
response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"]))
status = terminal_provider_status(response)
except Exception:
pass # Unreachable/expired is not proof of provider termination.
if status is not None or cancel:
self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time())
return status is not None
async def reconcile(self, worker_id=None):
"""Read terminal states for detached jobs and release only proven capacity."""
resolved = 0
async with self.recovery_lock:
for record in self._detached_records(worker_id):
resolved += await self._recover_provider_record(record)
await self._notify_capacity()
return resolved
async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False):
"""Explicit operator escape hatch for lost IDs/expired provider records.
This records a human's verification, not a cancellation or provider
status. It is never called automatically or allowed for live tasks.
"""
if provider_checked is not True:
raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it")
async with self.recovery_lock:
record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None)
if record is None:
raise ValueError("Unresolved inactive job not found for this worker")
self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(),
error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed")
await self._notify_capacity()
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
targets = {}
for rid, task in self.tasks.items():
run = self.runs[rid]
if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
continue
if execution_id and run["execution_id"] != execution_id:
continue
if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
continue
if not task.done():
# Repeated requests must not interrupt the first cancellation's cleanup.
if not task.cancelling():
task.cancel()
targets[rid] = task
if targets:
await asyncio.gather(*targets.values(), return_exceptions=True)
recovered = 0
async with self.recovery_lock:
for record in self._detached_records(worker_id, execution_id):
rid = record["record_id"]
# A task cancelled before its first step never enters _run's handlers.
# A failed start() may also leave a journalled but unstarted local slot.
if record.get("status") == "WAITING" and not record.get("job_id"):
run = self.runs.get(rid)
if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])):
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
run.pop("workflow", None)
run.pop("images", None)
run.pop("media", None)
await self._deliver(run, error="Serverless job cancelled before submission")
continue
# Do not send two cancellation requests for the same live job in
# one action. Its unresolved reservation survives for later recovery.
if rid in targets or collector_job_id or not record.get("job_id"):
continue
await self._recover_provider_record(record, cancel=True)
recovered += 1
await self._notify_capacity()
return len(targets) + recovered
def public_jobs(self, worker_id):
records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)]
# Never hide unresolved IDs behind the bounded completed history.
visible = [record for record in records if is_unresolved(record)]
visible.extend(record for record in records if not is_unresolved(record))
unresolved_count = sum(is_unresolved(record) for record in records)
return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")}
for record in visible[:max(20, unresolved_count)]]
+29
View File
@@ -0,0 +1,29 @@
"""Shared rules for provider capacity and durable recovery protection.
A local task ending is not evidence that provider work ended. RESOLVED is an
explicit operator acknowledgment, never a provider-reported state.
"""
PROVIDER_TERMINAL = frozenset({"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"})
TERMINAL = PROVIDER_TERMINAL | {"RESOLVED"}
PROVIDER_STATES = PROVIDER_TERMINAL | {"IN_QUEUE", "IN_PROGRESS"}
def terminal_provider_status(response):
status = response.get("status") if isinstance(response, dict) else None
return status if isinstance(status, str) and status in PROVIDER_TERMINAL else None
def is_unresolved(record):
return record.get("status") not in TERMINAL
def occupies_capacity(record):
# WAITING is journalled before admission and cannot have been submitted.
# Every other unresolved state is conservatively reserved, including a
# submission whose provider ID was lost and an unconfirmed cancellation.
return is_unresolved(record) and (record.get("status") != "WAITING" or bool(record.get("job_id")))
def belongs_to_pool(record, worker_id):
return str(record.get("worker_id")) == str(worker_id)
@@ -0,0 +1,225 @@
"""Bounded local native inputs for comfy-media-v1 requests.
Only stock LoadAudio/LoadVideo schemas are adapted. Original nodes and links
remain in the private request graph for endpoint validation; only the existing
image preparer rewrites filenames. No upload, transcode, resample or provider
operation belongs here. The caller still enforces the complete request-body cap.
"""
import hashlib
import io
import json
from pathlib import Path
from ...serverless_bridge import codec, contract, video_codec
from . import serverless_media
_LOADERS = {"LoadAudio": ("audio", "AUDIO"), "LoadVideo": ("file", "VIDEO")}
Error = serverless_media.ServerlessMediaError
def _loader_references(workflow):
if not isinstance(workflow, dict):
raise Error("Serverless workflow must be an API graph")
loaders = {}
for node_id, node in workflow.items():
if not isinstance(node, dict):
continue
class_type = node.get("class_type")
if not isinstance(class_type, str) or class_type not in _LOADERS:
continue
contract.identifier(node_id)
key, kind = _LOADERS[class_type]
inputs = node.get("inputs")
if not isinstance(inputs, dict) or set(inputs) != {key}:
raise Error(f"{class_type} inputs must contain only its native filename; extra controls are unsupported")
if not isinstance(inputs[key], str):
raise Error(f"{class_type} requires a literal filename; linked inputs are unsupported")
normalized, _suffix = serverless_media._normalize_media_reference(inputs[key])
loaders[node_id] = (kind, normalized)
references = []
used = set()
for node_id, node in workflow.items():
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
continue
for key, value in node["inputs"].items():
if not isinstance(value, (list, tuple)) or not value or not isinstance(value[0], str):
continue
if value[0] not in loaders:
continue
if type(value) is not list or len(value) != 2 or type(value[1]) is not int or value[1] != 0:
raise Error("Native media loader links require output slot 0")
contract.identifier(node_id)
contract.identifier(key)
references.append((node_id, key, value[0]))
used.add(value[0])
if used != set(loaders):
raise Error("Native media loader has no consumer")
return loaders, references
def _source_path(value, folder_paths):
"""Resolve native annotations; authorization/opening is owned by _read_path."""
annotation = serverless_media._MEDIA_ANNOTATION.search(value)
unannotated = value[:annotation.start()] if annotation else value
if Path(unannotated).is_absolute():
return unannotated
try:
return folder_paths.get_annotated_filepath(value)
except Exception:
raise Error("Native media input path is missing or invalid") from None
def _decode_audio(data, limit):
"""Match core nodes_audio.load, bounding each frame *before* accumulation.
Float32 output bytes are charged even for narrow integer PCM. Native float64
remains float64, as in f32_pcm (there is no silent dtype conversion). Raw
retained arrays and each live AV plane allocation have a separate <= limit
guard, so wider/padded raw frames cannot bypass the output budget. PyAV owns
the one frame being decoded; it is checked before to_ndarray or torch.cat.
Concatenation/conversion use bounded temporary copies, not an RSS guarantee.
"""
if limit <= 0:
raise Error("AUDIO decoded input exceeds byte limit")
try:
import av
import torch
from comfy_extras.nodes_audio import f32_pcm
except ImportError:
raise Error("Native AUDIO decoding is unavailable") from None
try:
with av.open(io.BytesIO(data)) as container:
if not container.streams.audio:
raise Error("Native AUDIO input has no audio stream")
stream = container.streams.audio[0]
rate = stream.codec_context.sample_rate
channels = stream.channels
if type(rate) is not int or rate <= 0 or type(channels) is not int or channels <= 0:
raise Error("Native AUDIO stream shape or sample rate is invalid")
frames = []
output_bytes = raw_bytes = 0
dtype = None
for frame in container.decode(streams=stream.index):
samples = frame.samples
if (type(samples) is not int or samples <= 0 or frame.sample_rate != rate
or len(frame.layout.channels) != channels):
raise Error("Native AUDIO frame shape or sample rate changed")
if frame.format.name not in {"s16", "s16p", "s32", "s32p", "flt", "fltp", "dbl", "dblp"}:
raise Error("Unsupported native AUDIO PCM format")
itemsize = frame.format.bytes
count = samples * channels
# Native f32_pcm promotes integer PCM to float32, not float64.
decoded_bytes = count * max(4, itemsize)
frame_bytes = count * itemsize
plane_bytes = sum(plane.buffer_size for plane in frame.planes)
if (decoded_bytes > limit - output_bytes or frame_bytes > limit - raw_bytes
or plane_bytes > limit - raw_bytes):
raise Error("AUDIO decoded input or raw frame exceeds byte limit")
output_bytes += decoded_bytes
raw_bytes += frame_bytes
array = frame.to_ndarray()
if array.ndim != 2 or array.size != count or array.nbytes != frame_bytes:
raise Error("Native AUDIO frame shape is invalid")
buf = torch.from_numpy(array)
if buf.dtype not in (torch.int16, torch.int32, torch.float32, torch.float64):
raise Error("Unsupported native AUDIO PCM dtype")
if dtype is not None and buf.dtype != dtype:
raise Error("Native AUDIO frame dtype changed")
dtype = buf.dtype
if buf.shape[0] != channels:
buf = buf.view(-1, channels).t()
if tuple(buf.shape) != (channels, samples):
raise Error("Native AUDIO frame shape is invalid")
frames.append(buf)
if not frames:
raise Error("Native AUDIO input has no decoded frames")
waveform = f32_pcm(torch.cat(frames, dim=1)).unsqueeze(0)
return {"waveform": waveform, "sample_rate": rate}
except Error:
raise
except Exception:
raise Error("Native AUDIO input could not be decoded safely") from None
def _artifact_id(kind, descriptor, payload):
digest = hashlib.sha256()
digest.update(kind.encode("ascii") + b"\0")
# Include interpretation (not just waveform bytes), notably sample_rate.
digest.update(json.dumps(descriptor, sort_keys=True, separators=(",", ":"), allow_nan=False).encode())
digest.update(b"\0")
digest.update(payload)
return "input-" + kind.lower() + "-" + digest.hexdigest()
async def prepare_media_inputs(workflow: dict) -> tuple[list, list, list]:
"""Return legacy images, inline typed artifacts, and every consumer reference.
Paths are read once into bounded bytes through the video adapter's no-follow
directory-fd reader. Decoders receive BytesIO, never a second path open.
Source files and unique decoded artifacts have independent aggregate caps.
Repeated content of the same kind is decoded once and charged once.
"""
loaders, references = _loader_references(workflow)
image_view = {node_id: node for node_id, node in workflow.items() if node_id not in loaders}
# Retain the strict scan for every other loader before expensive decoding.
serverless_media._iter_image_references(image_view)
if not loaders:
return await serverless_media.prepare_images(image_view), [], []
try:
import folder_paths
except ImportError:
raise Error("ComfyUI media directories are unavailable") from None
roots = list(serverless_media._folder_roots(folder_paths).values())
artifacts = {}
source_artifacts = {}
source_digests = set()
filenames = {}
loader_artifacts = {}
file_total = decoded_total = 0
for node_id, (kind, filename) in loaders.items():
# Avoid retaining file buffers merely to deduplicate repeat references.
if (kind, filename) in filenames:
loader_artifacts[node_id] = filenames[kind, filename]
continue
path = _source_path(filename, folder_paths)
data = video_codec._read_path(path, contract.MAX_ARTIFACT_BYTES, roots)
source_digest = hashlib.sha256(data).hexdigest()
if source_digest not in source_digests:
if len(data) > contract.MAX_TOTAL_BYTES - file_total:
raise Error("Native media source files exceed total byte limit")
file_total += len(data)
source_digests.add(source_digest)
source_key = (kind, source_digest)
if source_key not in source_artifacts:
if len(artifacts) >= contract.MAX_ARTIFACTS:
raise Error("Native media input artifact count exceeds limit")
limit = min(contract.MAX_ARTIFACT_BYTES, contract.MAX_TOTAL_BYTES - decoded_total)
if kind == "AUDIO":
value = _decode_audio(data, limit)
descriptor, payload = codec.encode_payload(kind, value, max_bytes=limit)
del value
else:
from comfy_api.latest import InputImpl
value = InputImpl.VideoFromFile(io.BytesIO(data))
descriptor, payload = video_codec.encode_video(value, max_bytes=limit, allowed_roots=roots)
del value
aid = _artifact_id(kind, descriptor, payload)
if aid not in artifacts:
if len(payload) > limit:
raise Error("Native media decoded inputs exceed total byte limit")
artifacts[aid] = contract.artifact_from_payload(
aid, kind, descriptor, payload, max_bytes=contract.MAX_ARTIFACT_BYTES)
decoded_total += len(payload)
source_artifacts[source_key] = aid
del payload
aid = source_artifacts[source_key]
loader_artifacts[node_id] = filenames[kind, filename] = aid
del data
refs = [{"node_id": node_id, "input_name": key, "artifact_id": loader_artifacts[source]}
for node_id, key, source in references]
# Shared node dictionaries deliberately propagate only legacy image rewrites.
images = await serverless_media.prepare_images(image_view)
return images, list(artifacts.values()), refs
@@ -0,0 +1,79 @@
"""Explicit media-v1 preflight and atomic, lossless result materialization."""
from urllib.parse import urljoin
import aiohttp
from ...serverless_bridge import contract
from ...serverless_bridge.contract import MAX_ARTIFACT_BYTES, PROFILE, validate_bindings
from .serverless_media import (
MAX_REDIRECTS, _CHUNK_SIZE, _DOWNLOAD_TIMEOUT_SECONDS,
_PublicResolver, _validate_artifact_url,
)
def validate_prepared_job(prepared):
"""Fail before journalling/submission; legacy jobs keep their original path."""
profile = prepared["worker"].get("media_profile", "legacy-image-v1")
media = prepared.get("media")
if profile == "legacy-image-v1" and media is None:
return False
if profile != PROFILE or media is None:
raise ValueError("Serverless media profile/request mismatch")
from .runpod_client import serialize_request
# The shared endpoint validator also checks exporter controls and input refs.
serialize_request(prepared["workflow"], prepared["images"],
prepared["worker"]["timeout_seconds"], media=media)
expected = validate_bindings(prepared["manifest"])
if expected != validate_bindings(media["outputs"]):
raise ValueError("Serverless media manifest identity mismatch")
if any(binding["participant_id"] != prepared["worker"]["id"] for binding in expected.values()):
raise ValueError("Serverless media participant identity mismatch")
return True
async def download_public_artifact(url, *, max_bytes=MAX_ARTIFACT_BYTES):
"""Contract-sized public HTTPS download, never using a provider session/key."""
if type(max_bytes) is not int or not 0 <= max_bytes <= MAX_ARTIFACT_BYTES:
raise ValueError("Media artifact exceeds the permitted size")
current_url = _validate_artifact_url(url)
connector = aiohttp.TCPConnector(resolver=_PublicResolver(), use_dns_cache=False, limit=4)
timeout = aiohttp.ClientTimeout(total=_DOWNLOAD_TIMEOUT_SECONDS, connect=10, sock_read=15)
try:
async with aiohttp.ClientSession(connector=connector, timeout=timeout, trust_env=False,
cookie_jar=aiohttp.DummyCookieJar()) as session:
for redirect_count in range(MAX_REDIRECTS + 1):
async with session.get(current_url, allow_redirects=False) as response:
if response.status in {301, 302, 303, 307, 308}:
location = response.headers.get("Location")
if not location or redirect_count >= MAX_REDIRECTS:
raise ValueError("Media artifact redirect is invalid")
current_url = _validate_artifact_url(urljoin(current_url, location))
continue
if response.status != 200:
raise ValueError("Media artifact download failed")
payload = bytearray()
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
if not isinstance(chunk, bytes) or len(payload) + len(chunk) > max_bytes:
raise ValueError("Media artifact exceeds the permitted size")
payload.extend(chunk)
return bytes(payload)
except Exception:
# Signed URLs and transport exception details must not reach the ledger.
raise ValueError("Media artifact download failed or exceeded its byte limit") from None
raise ValueError("Media artifact download failed")
async def collect_typed_results(output, media, manifest, *, download=None):
"""Decode every binding before returning any collector's native value lists."""
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
raise ValueError("Invalid serverless media output")
decoded = await contract.decode_response(output.get("media"), request_id=media["request_id"],
bindings=manifest,
download=download if download is not None else download_public_artifact)
grouped = {}
for binding in manifest:
values = decoded[binding["binding_id"]]
if not isinstance(values, list) or not values:
raise ValueError("Missing serverless media values")
grouped.setdefault(binding["job_id"], {}).setdefault(binding["socket"], []).extend(values)
return grouped
+283 -31
View File
@@ -1,4 +1,5 @@
import asyncio
import copy
import time
import uuid
@@ -13,7 +14,7 @@ from ..utils.constants import (
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
)
from ..utils.logging import debug_log, log
from ..utils.network import build_master_url
from ..utils.network import build_master_url, build_master_callback_url
from ..utils.trace_logger import trace_debug
from .schemas import parse_positive_float, parse_positive_int
from .orchestration.dispatch import (
@@ -21,7 +22,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,7 +35,11 @@ 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, _has_selected_serverless_workers
from ..image_list.execution_plan import prepare_image_list_execution
prompt_server = server.PromptServer.instance
@@ -74,6 +84,18 @@ def _resolve_enabled_workers(config, requested_ids=None):
elif not worker.get("enabled", False):
continue
if worker.get("type") == "serverless":
from ..utils.serverless_config import validate_serverless_worker
pool = validate_serverless_worker(worker)
if pool.get("media_profile") == "comfy-media-v1":
from ..serverless_bridge.contract import identifier
# Do not repair malformed media identities by stripping them.
identifier(pool["id"])
for slot in range(pool["jobs_per_run"]):
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
continue
raw_port = worker.get("port", worker.get("listen_port", 8188))
try:
port = int(raw_port or 8188)
@@ -131,19 +153,54 @@ def _is_load_balance_enabled(value):
def _prompt_requests_load_balance(prompt_index):
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
collector_ids = (prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("DistributedMediaCollector"))
for node_id in collector_ids:
inputs = prompt_index.inputs_by_node.get(node_id, {})
if _is_load_balance_enabled(inputs.get("load_balance", False)):
return True
return False
def _validate_media_workers(prompt_index, workers):
"""Reject incompatible selections before probing, pruning or dispatching.
Typed collectors currently have only a master/serverless transport. Check
all selected workers, not just reachable or load-balanced participants, so
a bad selection cannot silently turn into a different execution mode.
"""
typed = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
if typed and (prompt_index.nodes_for_class("DistributedCollector")
or prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")):
raise ValueError(
"DistributedMediaCollector cannot be mixed with DistributedCollector or "
"UltimateSDUpscaleDistributed in one prompt"
)
for worker in workers:
if worker.get("type") != "serverless":
if typed:
raise ValueError(
"DistributedMediaCollector does not support selected direct workers; "
"use master execution or comfy-media-v1 serverless pools"
)
continue
profile = worker.get("media_profile", "legacy-image-v1")
if profile not in ("legacy-image-v1", "comfy-media-v1"):
raise ValueError("Unsupported serverless media_profile")
if typed != (profile == "comfy-media-v1"):
raise ValueError(
"DistributedMediaCollector requires comfy-media-v1 serverless pools; "
"legacy collectors require legacy-image-v1"
)
async def _prepare_worker_payload(
worker,
prompt_index,
enabled_ids,
job_id_map,
master_url,
config,
delegate_master,
trace_execution_id,
worker_prep_semaphore,
@@ -152,22 +209,74 @@ async def _prepare_worker_payload(
):
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
async with worker_prep_semaphore:
_validate_media_workers(prompt_index, [worker])
worker_prompt = prompt_index.copy_prompt()
if worker.get("type") == "serverless":
from .orchestration.runpod_client import serialize_request
if worker.get("media_profile", "legacy-image-v1") == "comfy-media-v1":
from ..serverless_bridge.contract import PROFILE, VERSION
from .orchestration.serverless_media_prompt import compile_media_prompt
from .orchestration.serverless_typed_inputs import prepare_media_inputs
# The legacy pruner only recognizes Collector/Upscale roots.
# The typed compiler validates the complete private graph and
# prunes it itself, retaining every connected typed branch.
worker_prompt = apply_participant_overrides(
worker_prompt, worker["id"], enabled_ids, job_id_map,
"", delegate_master, prompt_index,
)
request_id = str(uuid.uuid4())
workflow, manifest = compile_media_prompt(worker_prompt, request_id)
images, artifacts, refs = await prepare_media_inputs(workflow)
media = {
"version": VERSION, "profile": PROFILE, "request_id": request_id,
"artifacts": artifacts, "inputs": refs, "outputs": copy.deepcopy(manifest),
}
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest,
"images": images, "media": media}
from .orchestration.serverless_prompt import compile_serverless_prompt
from .orchestration.serverless_media import prepare_images
worker_prompt = apply_participant_overrides(
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
job_id_map, "", delegate_master, prompt_index,
)
workflow, manifest = compile_serverless_prompt(worker_prompt)
images = await prepare_images(workflow)
# Fail during preparation, before regular workers or the master are
# dispatched. submit() uses this same serializer for the sent bytes.
serialize_request(workflow, images, worker["timeout_seconds"])
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images}
worker_master_url = build_master_callback_url(
worker,
config=config,
prompt_server_instance=prompt_server,
)
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"],
enabled_ids,
job_id_map,
master_url,
worker_master_url,
delegate_master,
prompt_index,
)
@@ -202,12 +311,14 @@ async def orchestrate_distributed_execution(
"""Core orchestration logic for the /distributed/queue endpoint.
Returns:
tuple[str, int]: (prompt_id, worker_count)
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
"""
ensure_distributed_state()
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
config = load_config()
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
(
@@ -217,8 +328,67 @@ async def orchestrate_distributed_execution(
media_sync_timeout_seconds,
) = _resolve_orchestration_limits(config)
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
for node in prompt_obj.values())
target_kwargs = {}
if image_list_mode:
if any(isinstance(node, dict) and node.get("class_type") in {
"DistributedCollector", "DistributedMediaCollector", "UltimateSDUpscaleDistributed"
} for node in prompt_obj.values()):
raise ValueError("Distributed Image List cannot be mixed with other distributed collectors")
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
target_kwargs["partial_execution_targets"] = outer_targets
prompt_index = PromptIndex(prompt_obj)
typed_media_mode = bool(prompt_index.nodes_for_class("DistributedMediaCollector"))
if typed_media_mode:
from ..serverless_bridge.contract import identifier
# PromptIndex's JSON copy would otherwise turn integer keys into strings.
for node_id in prompt_obj:
identifier(node_id)
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
image_list_mode = bool(image_list_start_ids)
if image_list_mode:
# Validate before probing or queueing anything so malformed boundaries and
# old worker releases fail before any upstream image is materialized.
other_distributed = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("DistributedMediaCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
)
if other_distributed:
raise ValueError(
"Distributed Image List v1 cannot be mixed with Distributed Collector, "
"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
# Pool identity is retained once, never expanded by Images per run.
# Validate profiles before reachability can hide an invalid selection.
workers += resolve_serverless_workers(config, requested_ids)
if not workers and delegate_master:
raise RuntimeError(
"Distributed Image List in orchestrator-only mode requires a selected direct worker "
"or comfy-media-v1 serverless pool"
)
else:
if typed_media_mode and requested_ids is not None:
configured_ids = {worker.get("id") for worker in config.get("workers", [])}
if any(worker_id not in configured_ids for worker_id in requested_ids):
raise ValueError("Unknown explicitly selected media worker ID")
workers = _resolve_enabled_workers(config, requested_ids)
_validate_media_workers(prompt_index, workers)
if typed_media_mode and delegate_master and not workers:
raise RuntimeError("Delegate-only DistributedMediaCollector requires a selected comfy-media-v1 pool")
load_balance_requested = _prompt_requests_load_balance(prompt_index)
trace_debug(
execution_trace_id,
@@ -234,26 +404,43 @@ async def orchestrate_distributed_execution(
),
)
# Respect master delegate-only configuration
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
if not workers and delegate_master:
# Keep the legacy Collector fallback. Image-list and typed-media delegate
# modes must not silently enable master participation.
if not image_list_mode and not workers and delegate_master:
trace_debug(
execution_trace_id,
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
)
delegate_master = False
active_workers, delegate_master = await select_active_workers(
active_workers, selected_delegate_master = await select_active_workers(
workers,
use_websocket,
delegate_master,
False if image_list_mode else delegate_master,
trace_execution_id=execution_trace_id,
probe_concurrency=worker_probe_concurrency,
)
if typed_media_mode and delegate_master and not active_workers:
raise RuntimeError("No selected comfy-media-v1 worker is reachable for delegate-only DistributedMediaCollector")
if image_list_mode:
from .orchestration.dispatch import preflight_image_list_workers
if not active_workers and delegate_master:
raise RuntimeError(
"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"]
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]
if not active_workers and delegate_master:
raise RuntimeError("No feature-compatible worker is available in orchestrator-only mode")
else:
delegate_master = selected_delegate_master
if load_balance_requested:
if load_balance_requested and not image_list_mode:
candidate_workers = list(active_workers)
if not delegate_master:
# Include master in load balancing only when master participation is enabled.
@@ -311,11 +498,19 @@ async def orchestrate_distributed_execution(
if not job_id_map:
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
return prompt_id, 0
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
queue_result = await queue_prompt_payload(
prompt_obj,
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
return (
queue_result["prompt_id"],
queue_result["number"],
0,
queue_result.get("node_errors", {}),
)
master_prompt = prompt_index.copy_prompt()
master_prompt = apply_participant_overrides(
@@ -326,10 +521,12 @@ async def orchestrate_distributed_execution(
master_url,
delegate_master,
prompt_index,
workflow_meta=workflow_meta,
)
if delegate_master:
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
if delegate_master and not image_list_mode:
collector_ids = (find_nodes_by_class(master_prompt, "DistributedCollector")
+ find_nodes_by_class(master_prompt, "DistributedMediaCollector"))
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
if upscale_nodes:
debug_log(
@@ -349,7 +546,7 @@ async def orchestrate_distributed_execution(
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
)
worker_payloads = []
if active_workers:
if active_workers and not image_list_mode:
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
worker_payloads = await asyncio.gather(
@@ -360,6 +557,7 @@ async def orchestrate_distributed_execution(
enabled_ids,
job_id_map,
master_url,
config,
delegate_master,
execution_trace_id,
worker_prep_semaphore,
@@ -370,7 +568,29 @@ async def orchestrate_distributed_execution(
]
)
if worker_payloads:
serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"]
regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"]
manager = None
if serverless_jobs:
from .orchestration.serverless_runtime import get_manager
# Repeat the manager's whole-batch identity/profile preflight *before*
# any direct dispatch, master queueing or collector queue allocation.
# Pure legacy jobs are a no-op in that validator; keep their path free
# of typed imports (and compatible with existing legacy media adapters).
for prepared in serverless_jobs:
if (prepared["worker"].get("media_profile", "legacy-image-v1") != "legacy-image-v1"
or prepared.get("media") is not None):
from .orchestration.serverless_typed_results import validate_prepared_job
validate_prepared_job(prepared)
manager = get_manager()
manager.ensure_capacity(len(serverless_jobs))
# Do not leave collector queues behind when offline compilation/media validation fails.
if not image_list_mode:
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
if regular_payloads:
await asyncio.gather(
*[
dispatch_worker_prompt(
@@ -381,13 +601,45 @@ async def orchestrate_distributed_execution(
use_websocket=use_websocket,
trace_execution_id=execution_trace_id,
)
for worker, wprompt in worker_payloads
for worker, wprompt in regular_payloads
]
)
prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
try:
queue_result = await queue_prompt_payload(
master_prompt,
workflow_meta,
client_id,
include_queue_metadata=True,
**target_kwargs,
)
except Exception:
if serverless_jobs:
async with prompt_server.distributed_jobs_lock:
for job_id in job_id_map.values():
prompt_server.distributed_pending_jobs.pop(job_id, None)
raise
prompt_id = queue_result["prompt_id"]
if manager is not None:
try:
await manager.start(execution_trace_id, prompt_id, serverless_jobs)
except Exception:
# Wake queued collectors with an explicit error if journal/capacity setup failed.
await manager.cancel(execution_id=execution_trace_id)
async with prompt_server.distributed_jobs_lock:
for prepared in serverless_jobs:
# Typed manifests have several sockets per collector, but
# each participant gets only one terminal collector event.
for job_id in {binding["job_id"] for binding in prepared["manifest"]}:
q = prompt_server.distributed_pending_jobs.get(job_id)
if q is not None:
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
raise
prompt_number = queue_result["number"]
node_errors = queue_result.get("node_errors", {})
trace_debug(
execution_trace_id,
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
)
return prompt_id, len(worker_payloads)
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
+98
View File
@@ -0,0 +1,98 @@
"""Master-only Runpod pool monitoring. Health checks never submit inference."""
from aiohttp import web
import server
from ..utils.config import load_config
from .orchestration.serverless_runtime import get_manager
from .orchestration.runpod_client import RunpodError
def _pool(worker_id):
return next((w for w in load_config().get("workers", [])
if w.get("id") == worker_id and w.get("type") == "serverless"), None)
async def _status(worker_id, fresh=False):
worker = _pool(worker_id)
if worker is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
jobs = []
try:
manager = get_manager()
jobs = manager.public_jobs(worker_id)
result = await manager.health(worker, fresh=fresh)
return web.json_response(dict(result, jobs=jobs))
except Exception as exc:
message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
@server.PromptServer.instance.routes.post("/distributed/serverless/test")
async def test_serverless_endpoint(request):
try:
data = await request.json()
worker_id = data.get("worker_id")
except Exception:
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
return await _status(worker_id, fresh=True)
@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
async def serverless_status_endpoint(request):
return await _status(request.match_info["worker_id"])
@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
async def cancel_serverless_endpoint(request):
try:
data = await request.json()
worker_id = data.get("worker_id")
except Exception:
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
if _pool(worker_id) is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
try:
manager = get_manager()
count = await manager.cancel(worker_id=worker_id)
return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
"message": "Cancellation requested for this master's jobs only"})
except Exception:
return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
async def _recover(request, *, acknowledge=False):
try:
data = await request.json()
worker_id = data.get("worker_id")
if not isinstance(worker_id, str) or not worker_id.strip():
raise ValueError()
record_id = data.get("record_id")
if acknowledge:
if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
raise ValueError()
except Exception:
return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
if _pool(worker_id) is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
try:
manager = get_manager()
if acknowledge:
await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
else:
result = {"resolved": await manager.reconcile(worker_id=worker_id)}
return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
except ValueError:
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this worker"}, status=409)
except Exception:
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
async def reconcile_serverless_endpoint(request):
return await _recover(request)
@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
async def resolve_serverless_endpoint(request):
return await _recover(request, acknowledge=True)
+33 -1
View File
@@ -16,12 +16,14 @@ from ..utils.logging import debug_log, log
from ..utils.network import (
build_worker_url,
get_client_session,
get_server_port,
handle_api_error,
normalize_host,
probe_worker,
)
from ..utils.constants import CHUNK_SIZE
from ..workers import get_worker_manager
from ..workers.ports import allocate_worker_ports
from .schemas import require_fields, validate_worker_id
from ..workers.detection import (
get_machine_id,
@@ -59,6 +61,13 @@ async def worker_ws_endpoint(request):
})
continue
if not isinstance(data, dict):
await ws.send_json({
"type": "dispatch_ack", "request_id": None, "ok": False,
"error": "Expected a websocket message object.",
})
continue
if data.get("type") != "dispatch_prompt":
await ws.send_json({
"type": "dispatch_ack",
@@ -79,10 +88,16 @@ async def worker_ws_endpoint(request):
continue
try:
target_kwargs = {}
if "partial_execution_targets" in data:
from ..utils.async_helpers import validate_execution_targets
target_kwargs["partial_execution_targets"] = validate_execution_targets(
prompt, data["partial_execution_targets"])
prompt_id = await queue_prompt_payload(
prompt,
workflow_meta=data.get("workflow"),
client_id=data.get("client_id"),
**target_kwargs,
)
await ws.send_json({
"type": "dispatch_ack",
@@ -277,6 +292,13 @@ def _get_cuda_info():
def _collect_network_info_sync():
"""Collect network/cuda info in a worker thread to avoid blocking route handlers."""
cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
master_port = get_server_port()
config = load_config()
worker_ports = []
if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
hostname = socket.gethostname()
all_ips = get_network_ips()
recommended_ip = get_recommended_ip(all_ips)
@@ -285,7 +307,9 @@ def _collect_network_info_sync():
"all_ips": all_ips,
"recommended_ip": recommended_ip,
"cuda_device": cuda_device,
"cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
"cuda_device_count": device_count,
"master_port": master_port,
"local_worker_ports": worker_ports,
}
@@ -449,6 +473,9 @@ async def launch_worker_endpoint(request):
if not worker:
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
if worker.get("type") == "serverless":
return await handle_api_error(request, "Serverless workers are managed by Runpod, not local processes", 400)
# Ensure consistent string ID
worker_id_str = worker_id
@@ -483,6 +510,8 @@ async def launch_worker_endpoint(request):
"message": f"Worker {worker['name']} launched",
"log_file": log_file
})
except ValueError as e:
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
except Exception as e:
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
@@ -541,6 +570,9 @@ async def get_local_worker_status_endpoint(request):
worker_statuses = {}
for worker in config.get("workers", []):
# Serverless pools have no host, port, or local process.
if worker.get("type") == "serverless":
continue
# Only check local workers
host = normalize_host(worker.get("host")) or ""
if not host or host in ["localhost", "127.0.0.1"]:
+1 -1
View File
@@ -1,7 +1,7 @@
# conftest.py — project-level pytest configuration.
#
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
# (from .distributed import ...) that fail when pytest tries to import it as a
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
# standalone module during Package.setup() for the root package node.
#
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
-51
View File
@@ -1,51 +0,0 @@
"""
ComfyUI-Distributed: thin entry point.
All implementation lives in workers/, nodes/, api/.
"""
import atexit
import os
import server
from .utils.config import ensure_config_exists
from .utils.logging import debug_log
from .utils.network import cleanup_client_session
from .workers import get_worker_manager
from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
from .upscale.job_store import ensure_tile_jobs_initialized
from .nodes import (
NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS,
ImageBatchDivider,
DistributedCollectorNode,
DistributedSeed,
DistributedModelName,
DistributedValue,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from . import api # noqa: F401 - triggers all @routes.* registrations
from .api.queue_orchestration import ensure_distributed_state
ensure_config_exists()
# Aiohttp session cleanup
async def _cleanup_session():
await cleanup_client_session()
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
# Initialize distributed job state on prompt_server
prompt_server = server.PromptServer.instance
ensure_distributed_state(prompt_server)
ensure_tile_jobs_initialized()
# Worker startup
if not os.environ.get('COMFYUI_IS_WORKER'):
atexit.register(sync_cleanup)
delayed_auto_launch()
register_async_signals()
+18 -8
View File
@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones.
- Dispatches the workflow to workers (`POST /prompt`).
- Probes and dispatches workers through `/distributed/worker_ws` by default.
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
- Queues the master workflow in ComfyUI’s prompt queue.
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
#### Fields
- `prompt` (required unless `workflow.prompt` is present, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
- A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
- This is not the normal visual workflow export from the ComfyUI editor.
- `workflow` (optional, object)
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- If you don’t care about UI metadata, you can omit it.
@@ -123,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
## Worker requirements (important)
For a worker to participate, it must be reachable from the master:
For a worker to participate, it must be reachable from the master. By default:
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
If `settings.websocket_orchestration=false`:
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
Also, for collector-based flows:
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
## Examples
### 1) Minimal `curl`
### 1) Minimal `curl` request envelope
```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-d @payload.json
```
Where `payload.json` contains at least:
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
```json
{
"prompt": {
"1": {"class_type": "KSampler", "inputs": {} }
"<node_id>": {
"class_type": "<node class>",
"inputs": {"<required_input>": "<value or connection>"}
}
},
"enabled_worker_ids": [],
"client_id": "external-client"
}
```
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
### 2) Python (`requests`)
```python
+38
View File
@@ -0,0 +1,38 @@
# Compressed Result Transfer
**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
| `transfer_format` | Type | Notes |
|---|---|---|
| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
Measured on a 121-frame 1280×720 batch (one worker's result):
| Format | Sent | vs legacy | PSNR |
|---|---|---|---|
| `legacy_png` | 446 MB | 1× | lossless |
| `png` | 163 MB | 2.7× | lossless |
| `ffv1` | 128 MB | 3.5× | lossless |
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
| `webp` q90 | 13 MB | 33× | 40.6 dB |
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
Notes:
- Choose the format on the master; the value is sent to workers with the workflow.
- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.
+180
View File
@@ -0,0 +1,180 @@
# Dynamic image-list queue
Process a list of images across the master, direct local or remote ComfyUI workers, and compatible Runpod serverless pools without assigning fixed chunks to each participant.
The master evaluates the input list once. By default, the master and each selected direct worker receive one image job at a time; a serverless pool contributes bounded concurrent slots. All slots share the same queue: when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. Each source frame is assigned once, and the master restores input order before running downstream nodes.
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
## Wiring
```text
Image list or IMAGE batch
|
Distributed Image List Start -------- context ---------+
| image |
v |
Per-image processing |
| IMAGE |
v v
Distributed Image List End
| image_list
v
Preview / Save / next stage
```
1. Install the same feature-capable version of Distributed on the master and every selected direct worker. Direct workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation. Serverless endpoints instead need the compatible media bridge and supervision described below.
2. Enable Distributed and select direct workers and/or **Typed media v1** serverless pools. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one eligible direct worker or compatible serverless pool. Serverless-only selection is supported. The master must be reachable from direct remote workers at its configured callback address; serverless jobs do not use this callback.
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; for each assigned item, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
## Inputs and outputs
Start accepts:
- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames; each frame becomes one job.
- Optional `mask`.
- Optional `value_1`, `value_2`, … inputs. Only `value_1` is shown at first; connecting it adds `value_2`, and so on, up to `value_10`. Each accepts an IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN.
`mask` and every value follow one rule:
- **1 value** is shared by every image.
- **One value per image** pairs them in order.
- **Any other count** is an error. Values are never padded or repeated silently.
Only `images` is split into frames. A batch connected to a value input stays a single value, so a 2-image batch can travel with each item (for example, two reference images per job). To pass a different batch to each image, connect a list with one batch per image.
Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`, then one `value_N` output per connected `value_N` input, with the same type. Unused value outputs are hidden. Seeds and strengths are ordinary values: connect a single number to share it, a list for one per image, or derive it from `item_index`.
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
End's optional `texts` is a connection-only socket. For workflows saved before this socket became connection-only, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
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
include the [standalone PR2 media bridge](../serverless_bridge/README.md), route
typed requests through its executor, and provide the required supervisor. Merely
installing the bridge nodes or choosing the profile is not sufficient. The
endpoint needs the workflow's models and custom nodes, but not the full
Distributed extension. A selected `legacy-image-v1` pool fails before the master
prompt is queued; it is not silently skipped in favor of local execution.
Provider health/authentication probing does not verify bridge/model installation.
- Each pool is selected once by its configured pool ID. **Images per run
(`jobs_per_run`) does not multiply input images or set image-list parallelism.**
An assigned cloud job processes exactly one source item and its paired values.
- Cloud slots are bounded by the existing **16 outstanding jobs per pool** cap,
available manager capacity (**64 pending jobs** across pools), reservations,
and the source item count. Unresolved submissions and unconfirmed cancellations
still consume capacity. Selecting more items does not enqueue all of them as
provider jobs at once. Endpoint-side worker limits can reduce actual concurrency.
- Cloud slots, direct workers and the optional master pull from the same queue;
no region is cloned once per Images-per-run setting. Results are restored to
source order even when participants finish out of order.
- 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 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
finite; connected masks must match their image grid. Unsupported values fail
rather than being resized, flattened or converted to lossy images.
- Inputs are inline and the **complete serialized Runpod request is capped at
10,000,000 bytes**, including workflow and base64. Default inline result JSON
has the same cap. The PR2 contract's larger artifact ceilings do not remove
these inline limits; larger outputs need an operator-provided artifact publisher.
There is no built-in large-input uploader. See [typed media limits](serverless-media.md#size-and-operational-limits).
- Missing/duplicate/mismatched results or provider failures fail the region.
Cancellation closes only this region's owned cloud jobs and collector queues,
including jobs whose dispatch was cancelled before returning a handle. It does
not clear unrelated pool work. No failed or ambiguous job is retried automatically.
- A provider cancellation or timeout is **not proof that endpoint ComfyUI stopped**.
The PR2 supervisor must quarantine/reconcile or terminate an uncertain endpoint
before reuse. Preserve its retained files until execution and readers are known
to have stopped; see the [bridge uncertainty policy](../serverless_bridge/README.md).
## Scope and failure behavior
- This version supports master, direct workers and `comfy-media-v1` serverless pools in one image-list region. Direct capability/idle HTTP probes are never sent to serverless pools; pool probing uses the existing serverless health/authentication path.
- Worker eligibility is determined before execution and checked again before dispatch. Direct workers must be idle in orchestrator-only mode; master-participating mode can use reachable compatible peers without requiring an empty queue. Transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
- Omitted worker IDs select enabled direct workers and compatible serverless pools. Explicit IDs select the named configured workers/pools, including disabled ones; an empty list stays empty and unknown IDs fail. Pool IDs are retained in Start's hidden selection metadata, not expanded participant IDs. Local direct workers use the master's loopback callback; remote direct workers use its configured public address.
- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no workers are selected or reachable. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
- 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 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.
- Image-list ordering/context state is in memory. Restarting the master does not resume an interrupted region. Cloud job IDs and uncertainty state remain in the existing durable serverless ledger for reconciliation, not automatic list resumption. Abandoned master contexts and expired direct-worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
### Time and size limits
`settings.image_list_job_timeout_seconds` defaults to `0` (no arbitrary total-run cap). A positive finite number opts into an overall deadline with owned-work cleanup; negative, non-finite or malformed settings fail explicitly. There is no fixed one-hour region deadline. The setting is separate from the existing worker liveness timeout.
Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each.
**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/mask/value combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
## Verification
The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started.
Offline cloud integration tests additionally cover mixed-role and cloud-only
scheduling, out-of-order image/mask/text pairing, one job per source item regardless
of Images per run, profile/selection rejection before queueing, exact cleanup on
setup/dispatch/result failures, and the real CPU ComfyUI `processing_interrupted`
signal while awaiting a mocked cloud result. Provider calls are mocked: this is
not live Runpod, endpoint-supervisor or GPU-workflow acceptance.
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
bash scripts/test-web.sh
```
Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.
+158
View File
@@ -0,0 +1,158 @@
# Typed serverless media (PR2)
`comfy-media-v1` adds typed list results, exact by default, for `IMAGE`, `AUDIO`, `MASK`,
`STRING`, and native ComfyUI `VIDEO`. It is an **explicit endpoint profile**, not
an automatically detected capability. Omitting `media_profile` still selects
`legacy-image-v1`; existing image collectors and handlers keep their old contract.
See [legacy serverless setup](serverless-workers.md) for credentials and recovery.
## Use a typed collector
1. Prepare an endpoint with the [standalone media bridge](../serverless_bridge/README.md)
and the workflow's models/custom nodes. The endpoint needs only the bridge,
**not the full ComfyUI-Distributed extension**.
2. In the serverless worker settings, set **Endpoint media profile** to
**Typed media v1 (requires endpoint bridge)** (`media_profile: "comfy-media-v1"`).
Start with **Images per run = 1**. This existing setting counts remote jobs,
not the number of typed items each job returns.
3. Use **Distributed Media Collector** (`DistributedMediaCollector`), connect at
least one typed input, and connect its outputs to the appropriate downstream
nodes. For literal text, connect a `PrimitiveString` node to `texts`.
4. Queue through Distributed. **Delegate Only** excludes local upstream
generation; normal collection includes the master's local values.
**Test Saved Endpoint** checks provider authentication/health, not bridge
installation, model availability, or typed execution. Selecting a profile does
not install or upgrade anything on the endpoint.
The collector has five optional inputs and five **list outputs**, in this order:
| Socket / output index | Native type | Each list item |
| --- | --- | --- |
| `images` / 0 | `IMAGE` | One BHWC tensor, including its batch dimension |
| `audio` / 1 | `AUDIO` | `{ "waveform": BCT_tensor, "sample_rate": positive_int }` |
| `masks` / 2 | `MASK` | One BHW tensor |
| `texts` / 3 | `STRING` | Exact text, including whitespace or an empty string |
| `video` / 4 | `VIDEO` | One supported native video object |
For each connected socket, order is **master first, then the prepared participant
order, then that participant's item order**—never completion order. Delegate-only
mode omits the master. An unconnected socket returns `[]`. A connected remote
socket must return 1–64 items; an empty string is an item, an empty list is not.
Lists are not concatenated into one tensor, resampled, resized, or aligned across
sockets. Different items may have different valid shapes or audio sample rates.
The entire remote response is validated and decoded before any of its results
are delivered; missing, duplicate, mismatched, or corrupt results fail collection.
## Image return format and quality
`transfer_format` defaults to **`exact`** and `transfer_quality` to **90**.
These controls affect **remote IMAGE returns only**, not master-local images,
uploaded inputs, MASK, AUDIO, STRING or native VIDEO values.
- `exact`: preserve floating-point dtype and values, with no image quantization.
- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float
tensors. Image values are clipped to [0,1] and quantized as in the existing
Collector. `legacy_png` uses PNG without compression; `png` compresses it.
- `jpeg`, `webp`: lossy still-image compression.
- `h264`, `h265`, `av1`: lossy video encoding of each IMAGE tensor's frame batch.
This does not transcode the native `video` socket or combine separate list items.
- Quality 1–100 affects lossy formats only, using the existing Collector's
quality-to-CRF mapping for video. Dimensions, frame order and channel count
are preserved; alpha is carried separately where needed, not silently dropped.
Both master and endpoint must have this bridge revision for compressed image
returns. Pillow handles still formats and CPU/software PyAV encoders handle video;
unavailable codecs fail clearly rather than silently changing the selection.
Compression can reduce result egress and inline response size, but does not relax
encoded/decoded byte budgets or the input request limit. Compressed IMAGE batches
are capped at 240 frames and 8192 pixels per dimension, with a conservative
256 MiB aggregate decode budget (including codec scratch and alpha). The original Collector's
settings, default and fallback behavior are unchanged.
## Inputs and fidelity
Automatic local-file transfer supports stock `LoadImage`, plus explicit stock
`LoadAudio` (`audio` filename) and `LoadVideo` (`file` filename) adapters. Audio and
video loaders require a literal filename, no extra controls, and output slot 0.
Files must be under ComfyUI's input/output/temp roots. Custom file loaders are not
supported; detected unsupported file references fail rather than silently
becoming endpoint-local paths. Models and custom nodes are never transferred.
- `LoadImage` uses the existing bounded file/base64 image staging path.
- `LoadAudio` is decoded on the master using native PCM semantics with bounded
accumulation, then transferred as a waveform and sample rate. It does not
resample; integer PCM conversion follows ComfyUI's `f32_pcm` behavior.
- `LoadVideo` transfers opaque file bytes as native `VideoFromFile`, without
decoding frames or transcoding the container.
- Typed tensor transport is `tensor-le-v1`: contiguous little-endian float16,
bfloat16, float32, or float64 bytes with exact dtype/shape/layout metadata.
Only dense, finite tensors with positive dimensions are supported. IMAGE has
1, 3, or 4 channels. Bridge export explicitly detaches/stages tensors to CPU;
the pure codec itself rejects non-CPU storage. Remote decoding returns
independently owned CPU values. STRING transport is exact UTF-8 (`utf8-v1`).
### Native VIDEO is not a generic video converter
The adapter supports the pinned native classes **exactly**, not arbitrary
subclasses or objects exposing `save_to()` / `get_components()`:
- `VideoFromComponents`: images, optional matching BHW or BHW1 alpha, optional
audio, positive rational frame rate, bounded JSON metadata, positive integer
bit depth, and `sRGB`, `HDR`, or `HDR PQ` color space are preserved.
- `VideoFromFile`: regular local files within allowed media roots or an exact
`BytesIO`; opaque bytes, finite signed start/duration, and optional four-integer
crop are preserved without normalization.
- `VideoFromList`: 1–64 flat exact `VideoFromFile` segments, plus optional
`complete_audio`, in order.
Unknown classes, incompatible native state, nested/component segment lists,
unsupported tensors, invalid metadata, and oversized payloads are rejected.
There is no fallback transcode, frame extraction, color conversion, or container
repair/validation. A downstream native consumer still needs to read the original
container. Decoded file-backed VIDEO values own independent `BytesIO` storage;
component tensors also own their storage, so staging-file cleanup cannot
invalidate a returned VIDEO.
## Size and operational limits
- The typed contract caps each artifact at **256 MiB**, aggregate artifact bytes
at **512 MiB**, and each artifact list at **128** entries. These are validation
ceilings, not a promise that inline requests can carry that much data.
- The built-in master packs inputs inline and enforces **10,000,000 bytes for the
complete serialized Runpod `/run` request**, including workflow, policy and
base64. Legacy image reads also retain their 10,000,000-byte aggregate guard.
There is **no built-in master-side large-input uploader** in PR2.
- The endpoint defaults to inline results and caps its result JSON at
**10,000,000 bytes**. Larger results require an operator-supplied publisher
returning expiring public HTTPS artifact URLs. See the bridge guide for
bounded download/upload responsibilities and retention.
- Remote typed transport is **serverless only**. Direct-worker typed transport
is not implemented. Do not mix a typed collector with legacy image collectors,
unsupported distributed/batch nodes, or collector-to-collector dependencies.
Downstream processing stays on the master. Image-list scheduling is separate
PR3 work, not part of this contract.
## Timeout is not proof that Comfy stopped
Provider cancellation, a handler timeout, or cancellation of its Python task
**does not establish that the local Comfy prompt stopped**. The bridge never
posts `/interrupt`: the pinned core route checks a queue snapshot and then
interrupts globally, so even a prompt ID does not make it race-free.
If submission/termination or callback teardown cannot be confirmed, the bridge
raises `MediaExecutionUncertain` and retains owned input/exporter files—even
after terminal Comfy history if a downloader/publisher ignores cancellation.
Callback teardown has a separate 10-second grace; late callback exit does not
clean retained files. Callbacks must be cooperative async code; synchronous
event-loop blocking cannot be hard-timed-out.
The endpoint operator must **quarantine the owning worker/environment, reconcile
the local prompt and readers or terminate that environment, and establish that
they have stopped before cleanup or worker reuse**. Do not automatically retry
the request or turn every handler exception into a safely reusable terminal
provider job. Master-side provider status/reconciliation alone is not proof
that a surviving local Comfy process or callback has stopped.
These docs describe the implementation and its offline tests; they do not claim
live Runpod/provider or GPU end-to-end validation.
+173
View File
@@ -0,0 +1,173 @@
# Runpod serverless workers (preview)
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
normal ComfyUI API workflow per requested image and collects the returned images.
The endpoint workers do not need ComfyUI-Distributed installed.
## Setup
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
2. Enter its name, Runpod **endpoint ID** (not a URL; a pasted `https://api.runpod.ai/v2/<id>/…`
URL is trimmed to the ID), and API key.
3. Start with **Images per run = 1** and **Timeout = 1800**.
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
endpoint health only; it does not start GPU inference or prove model/node compatibility.
5. Enable the worker when you want queued Distributed workflows to use it.
6. Use an image **Distributed Collector** and queue through Distributed as usual.
Enable **Delegate Only** if the master should collect without generating locally.
Keys are write-only. Leave the key field blank to retain a saved key, or select
**Clear saved API key** to remove it. The checkbox is unavailable while the
worker is enabled, so turn the worker off first. A saved key is required before
enabling a worker. Use HTTPS or a trusted local connection when
entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
## Endpoint contract
The endpoint must accept this generic input shape; named application modes are
not used:
```json
{
"input": {
"workflow": {
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}}
},
"images": []
}
}
```
For image inputs, `images` contains `{ "name": "input.png", "image": "<base64>" }`
entries and the handler must rewrite `LoadImage` references to the staged files.
The adapter transfers local input images; it does not copy models or install nodes.
All referenced models and custom nodes must already exist in the endpoint image.
On completion, Runpod's `output.images` must contain artifact records such as:
```json
{
"filename": "example_00001_.png",
"subfolder": "",
"artifact_type": "output",
"artifact_path": "example_00001_.png",
"type": "base64",
"data": "<base64>"
}
```
Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"`
with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without
Runpod credentials; private-network destinations and unsafe redirects are refused.
The handler must preserve injected `SaveImage.filename_prefix` values so results
can be matched to the right collector and slot. Duplicate basenames in different
subfolders remain different artifact identities.
A Runpod job marked `COMPLETED` is still rejected if `output.error` or
`output.errors` is nonempty, required collector outputs are missing, or image
artifacts are invalid. Failures are surfaced rather than silently returning an
incomplete serverless batch.
## Scheduling and cancellation
- **Images per run** (1–16) is how many separate Runpod jobs this worker submits
for each queued run. Each returns one image and receives its own Distributed
Seed/Value overrides. All of them are submitted at once.
- How many of those jobs actually run in parallel is decided by the Runpod
endpoint's **max workers**, not by Distributed. Jobs beyond that wait in
Runpod's queue. To run every image at the same time, keep Images per run at or
below the endpoint's max workers.
- Distributed also keeps a fixed limit of 16 outstanding jobs per serverless
worker (and 64 across all of them). Retained jobs after a restart, ambiguous
submissions, and unconfirmed cancellations keep counting against it until
reconciled. Other applications using the same endpoint are unaffected.
- Older configs may contain a `concurrency` field from the previous separate
setting. It is still accepted when loading or saving, but ignored and dropped.
- **Load balance** selects one participant, so it sends one job instead of
Images per run.
- **Timeout** includes Runpod queue/cold-start time, execution, and output
retrieval. Provider policy limits are submitted as well.
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
- Interrupting a collector or using the panel cancellation action requests
cancellation only for jobs submitted by this master. It never purges the endpoint.
- A network failure during submission is not retried automatically: the job
may already be billable. Check the Runpod job list before queuing again.
## Persistence and recovery
Credentials live outside `gpu_config.json` and Git in the master ComfyUI user
folder under `distributed/serverless/credentials.json`, with restrictive POSIX
permissions. They are omitted from public configuration, workflow exports, and
browser configuration state. The first version requires a POSIX (Linux/macOS)
master for serverless secret storage; existing non-serverless workers remain
available on Windows.
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
not image data, workflow payloads, or API keys. A restarted master cannot resume
an interrupted collector. It does not automatically resubmit jobs. Retained
IDs are available through `/distributed/serverless/status/<worker_id>`; the worker's
cancel action can request cancellation of known pre-restart jobs. If no job ID
was received, check Runpod manually. Cancellation is best-effort and may race
with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation.
Unresolved records protect the worker's endpoint, API key, worker type,
and deletion even when no local task remains. Disabling the worker is still allowed;
it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are
closed locally, while interrupted `SUBMITTING` records without IDs become
`UNKNOWN` and remain protected.
### Reconcile or acknowledge recovery
Recovery uses the same trusted-master API as the panel; it never resubmits a job.
1. Read `GET /distributed/serverless/status/<worker_id>`. The `jobs` list includes
every unresolved record with its local `record_id`, provider `job_id` (if
received), and state. These records remain visible if endpoint health fails.
2. Request cancellation from the panel, or send
`POST /distributed/serverless/reconcile` with the body below to read the
retained provider IDs. Only a confirmed terminal state releases protection.
Network errors, expired IDs, or ambiguous responses do not.
```json
{"worker_id": "YOUR_WORKER_ID"}
```
3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
manually first. **Only after verifying that the job is no longer running or
queued**, send `POST /distributed/serverless/resolve` for that one record:
```json
{"worker_id": "YOUR_WORKER_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
```
This records `RESOLVED`: an explicit operator acknowledgment, not a
provider-confirmed cancellation. It cannot resolve a live local task and does
not stop remote compute. It releases the record's reservation and may let
already-waiting jobs proceed. Never automate this acknowledgment as a retry
mechanism; an unverified job can still be billable.
Worker setting changes become available once all its local tasks and unresolved records
are finished. No credentials or ledger files need to be edited by hand.
## First-version limits
- Image collectors only. Audio/video transfer and audio collectors are rejected.
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
node types are rejected before submission.
- Downstream post-collector processing remains local; the remote graph contains
only the upstream compute required for image saves.
- Inputs must be supported images under ComfyUI's input/output/temp roots.
- Raw image input reads are guarded at 10,000,000 bytes total, and the complete
serialized `/run` JSON request body (workflow, base64 images, and policy) must
also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image
budget is lower and depends on workflow size. Large-media/object-storage upload
support is not included.
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
scaling changes, or inbound master callback/tunnel is needed.
This feature is verified with offline provider/transport tests. A successful
health test is not a paid end-to-end generation test; test your deployed handler
and model inventory with one image per run before increasing it.
+53
View File
@@ -0,0 +1,53 @@
# Native V3 node API
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
that API; there is no V1 registration fallback.
## Compatibility
- All eight node IDs, display names, categories, visible input order, defaults
and output order are retained. Existing execution algorithms remain in the
private collector, utilities and upscale modules.
- The image/audio dividers explicitly declare ten typed outputs. Their existing
frontend extensions still show only the selected number of outputs. This
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
IMAGE/AUDIO socket indices or the ten returned values.
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
its existing prompt-input names and Python defaults via `accept_all_inputs`;
it does not become visible widgets.
- The collector retains list-input handling. Upscale retains its always-changing
fingerprint and creates a private runtime helper for each V3 execution, so
mutable helper state is not attached to sanitized V3 class clones.
- Routes, distributed state and the existing worker startup/shutdown hooks are
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
Worker mode still suppresses automatic worker launch. `distributed.py` is no
longer an entrypoint.
## Verification
Run ordinary unit tests from this repository:
```bash
python -m pytest tests -q -o addopts=
```
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
```
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
all node schemas with the V1 fixture, checks all five bundled workflows' node
IDs and socket/link contracts, exercises injected worker metadata, image/audio
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
and decodes an actual preview PNG referenced by executor history.
The upscale GPU/model boundary is mocked to check argument forwarding and
per-execution helper isolation. Full checkpoint inference, browser canvas
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
workers or model downloads are started; preview files use a temporary scratch
directory. The test does not install this branch into a live custom-node folder.
+4 -2
View File
@@ -74,12 +74,14 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
> ComfyUI instances running on a cloud service like Runpod.
For faster result transfers and lower bandwidth usage, see [Compressed Result Transfer](compressed-result-transfer.md).
### Deploy Cloud Worker on Runpod
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
**On Runpod:**
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
@@ -92,7 +94,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
- SAGE_ATTENTION: optional optimisation (set to true/false)
5. Deploy your pod.
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
```
# Download from CivitAI
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
+1
View File
@@ -0,0 +1 @@
"""Dynamic per-image image-list execution internals."""
+70
View File
@@ -0,0 +1,70 @@
"""Prepare image-list execution without changing the saved/API prompt or registry."""
from __future__ import annotations
import copy
import sys
from .runtime import (
END_CLASS,
_ancestors_for_inputs,
_downstream_until,
_image_list_pairs,
_is_output_node,
_prompt_links,
validate_image_list_regions,
)
def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]:
visited = set()
pending = list(roots)
while pending:
node_id = pending.pop()
if node_id in visited:
continue
visited.add(node_id)
pending.extend(edges[node_id] - visited)
return visited
def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
"""Return an execution copy and its outer-master output targets.
Disposable per-item side branches are removed before boundary validation.
Required output-node utilities remain normal dependencies, but must not be
independent outer roots. Callers must use these targets (and End-only targets
for item execution); this does not suppress side effects inside required nodes.
"""
pairs = _image_list_pairs(prompt)
prepared = copy.deepcopy(prompt)
upstream = {node_id: set() for node_id in prepared}
downstream = {node_id: set() for node_id in prepared}
for node_id, node in prepared.items():
for _, source, _ in _prompt_links(node, prepared):
upstream[node_id].add(source)
downstream[source].add(node_id)
# Never erase another region or its dependencies to make a bad graph pass.
boundaries = {node_id for pair in pairs for node_id in pair}
required = _reachable(boundaries, upstream)
discarded = set()
for start_id, end_id in pairs:
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, ("images", "masks", "texts"))
outside = _downstream_until(prepared, start_id, end_id) - region_nodes
# Preserve bypass joins and their whole chain into post-End computation.
# They must fail boundary validation, not disappear with debug previews.
post_end = _reachable({end_id}, downstream)
protected = required | _reachable(post_end, upstream)
discarded.update(outside - protected)
for node_id in discarded:
del prepared[node_id]
regions = validate_image_list_regions(prepared, explicit_targets=True)
per_item = set().union(*(region["branch_ids"] for region in regions))
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
targets = [
node_id for node_id, node in prepared.items()
if node_id not in per_item and isinstance(node, dict)
and (node.get("class_type") == END_CLASS or _is_output_node(node, registry))
]
return prepared, targets
+174
View File
@@ -0,0 +1,174 @@
"""Server-owned background cleanup for image-list runtime state."""
from __future__ import annotations
import asyncio
import logging
import math
import time
from collections.abc import Iterable
from typing import Any, cast
LOGGER = logging.getLogger(__name__)
DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0
DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
_LIFECYCLE_MARKER = object()
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
"""Prove outer roots from the running queue, not a forgeable prompt marker.
Plain /prompt remains strict when ComfyUI selected all outputs, including a
flagged per-item dependency. No saved workflow field grants this exemption.
"""
if not prompt_id:
return False
from .execution_plan import prepare_image_list_execution
prepared, targets = prepare_image_list_execution(prompt)
if prepared != prompt or not targets:
return False
queue = getattr(prompt_server, "prompt_queue", None)
for name in ("get_current_queue_volatile", "get_current_queue"):
reader = getattr(queue, name, None)
if not callable(reader):
continue
try:
running, _pending = reader()
for entry in running:
if (isinstance(entry, (tuple, list)) and len(entry) > 4
and entry[1] == prompt_id and entry[2] == prompt):
roots = entry[4]
return (isinstance(roots, (list, tuple))
and all(isinstance(root, str) for root in roots)
and len(roots) == len(targets) and set(roots) == set(targets))
return False
except Exception:
continue
return False
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
queue = getattr(prompt_server, "prompt_queue", None)
readers = (
getattr(queue, "get_current_queue_volatile", None),
getattr(queue, "get_current_queue", None),
)
errors: list[Exception] = []
for reader in readers:
if not callable(reader):
continue
try:
snapshot = reader()
if not isinstance(snapshot, (tuple, list)) or len(snapshot) < 2:
raise TypeError("prompt queue snapshot must contain running and pending entries")
running, pending = snapshot[0], snapshot[1]
prompt_ids: set[str] = set()
for entries in (running, pending):
if not isinstance(entries, Iterable):
raise TypeError("prompt queue entries must be iterable")
for entry in entries:
if isinstance(entry, (tuple, list)) and len(entry) > 1:
prompt_id = entry[1]
if isinstance(prompt_id, str):
prompt_ids.add(prompt_id)
return prompt_ids
except Exception as exc:
errors.append(exc)
if errors:
raise RuntimeError("unable to read the prompt queue") from errors[-1]
raise RuntimeError("prompt queue has no supported queue reader")
def _sweep_orphan_contexts(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
queue_read_ok = True
try:
queued_prompt_ids = _queue_prompt_ids(prompt_server)
except Exception:
# A failed queue read must not make a submitted context look orphaned.
LOGGER.warning("Could not read prompt queue during image-list lifecycle sweep", exc_info=True)
queued_prompt_ids = set()
queue_read_ok = False
now = time.monotonic()
with runtime._LOCK:
for context_id, state in list(runtime._MASTER_JOBS.items()):
if getattr(state, "active", False):
continue
prompt_id = getattr(state, "prompt_id", None)
if prompt_id is not None:
if not queue_read_ok or prompt_id in queued_prompt_ids:
continue
created_at = getattr(state, "created_at", None)
expired = (
isinstance(created_at, (int, float))
and not isinstance(created_at, bool)
and math.isfinite(created_at)
and now - created_at >= orphan_grace_seconds
)
if expired and runtime._MASTER_JOBS.get(context_id) is state:
runtime.remove_master_job(context_id)
def _sweep_once(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
try:
runtime._sweep_worker_items()
except Exception:
LOGGER.exception("Could not sweep expired image-list worker staging")
_sweep_orphan_contexts(prompt_server, runtime, orphan_grace_seconds)
async def _sweep_loop(
prompt_server: Any,
runtime: Any,
interval_seconds: float,
orphan_grace_seconds: float,
) -> None:
while True:
try:
_sweep_once(prompt_server, runtime, orphan_grace_seconds)
except Exception:
# Keep the server-owned task alive if one pass encounters unexpected state.
LOGGER.exception("Image-list lifecycle sweep failed")
await asyncio.sleep(interval_seconds)
def install_lifecycle(
prompt_server: Any,
runtime: Any,
*,
interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS,
orphan_grace_seconds: float = DEFAULT_ORPHAN_GRACE_SECONDS,
) -> None:
"""Register one cancellable periodic sweep in the prompt server's app context."""
interval_seconds = float(interval_seconds)
orphan_grace_seconds = float(orphan_grace_seconds)
if not math.isfinite(interval_seconds) or interval_seconds <= 0:
raise ValueError("interval_seconds must be a positive finite value")
if not math.isfinite(orphan_grace_seconds) or orphan_grace_seconds < 0:
raise ValueError("orphan_grace_seconds must be a non-negative finite value")
app = getattr(prompt_server, "app", None)
cleanup_ctx = getattr(app, "cleanup_ctx", None)
if cleanup_ctx is None:
raise TypeError("prompt_server must expose an app with cleanup_ctx")
if any(getattr(factory, "_image_list_lifecycle_marker", None) is _LIFECYCLE_MARKER
for factory in cleanup_ctx):
return
async def lifecycle_context(_app):
task = asyncio.create_task(
_sweep_loop(prompt_server, runtime, interval_seconds, orphan_grace_seconds),
name="image-list-lifecycle-sweeper",
)
try:
yield
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
cast(Any, lifecycle_context)._image_list_lifecycle_marker = _LIFECYCLE_MARKER
cleanup_ctx.append(lifecycle_context)
+184
View File
@@ -0,0 +1,184 @@
"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.
Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
process interrupt, resets progress handlers, changes RAM-cache callbacks and
starts/ends a model tracker. Those belong to the suspended parent executor.
Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
with private provider-free caches and a private server facade. Model outputs are
registered with the existing parent tracker, whose lifetime includes this work.
Core progress and GraphBuilder defaults are process-global. Borrow them only
inside the owned execution thread, restore them in finally, and join that thread
before End returns. Cancellation never clears or raises the global interrupt:
it stops locally at node/progress boundaries. A non-cooperative synchronous node
must finish before cleanup can return; abandoning its thread is not safe.
"""
from __future__ import annotations
import asyncio
from concurrent.futures import ThreadPoolExecutor
import threading
from types import SimpleNamespace
from .runtime import ImageListCancelledError
_EXECUTION_LOCK = threading.Lock()
class LocalExecutionSlot:
def __init__(self, *, model_callback=None):
self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
self._stop = threading.Event()
self._future = None
self._caches = None
self._model_callback = model_callback
self._closed = False
self._loop = None
self._task = None
self._state_lock = threading.Lock()
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
if self._closed or self._stop.is_set():
raise ImageListCancelledError("Local image-list slot has stopped")
if self._future is not None and not self._future.done():
raise RuntimeError("Local image-list slot already has an in-flight item")
self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
return self._future
def request_stop(self):
self._stop.set()
# Wakes a graph waiting on an async node as well as stopping synchronous
# nodes at their next progress/between-node checkpoint.
with self._state_lock:
if self._loop is not None and self._task is not None:
self._loop.call_soon_threadsafe(self._task.cancel)
async def close(self):
self.request_stop()
try:
if self._future is not None:
# Do not cancel this Future: cancelling an asyncio wrapper cannot
# kill its running thread, and falsely makes cleanup look done.
waiter = asyncio.wrap_future(self._future)
while not waiter.done():
try:
await asyncio.shield(waiter)
except asyncio.CancelledError:
continue
except Exception:
break # The scheduler propagates execution errors.
if waiter.done() and not waiter.cancelled():
waiter.exception()
finally:
self._pool.shutdown(wait=True, cancel_futures=True)
self._caches = None
self._closed = True
def _run(self, prompt, prompt_id, end_id, workflow_meta):
# The normal Comfy prompt worker is single-threaded and blocked in End.
# Fail closed if another local slot is nevertheless executing, rather
# than interleave writes to core's process-global execution state.
if not _EXECUTION_LOCK.acquire(blocking=False):
raise RuntimeError("Another local image-list executor is already active")
try:
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
except asyncio.CancelledError as exc:
raise ImageListCancelledError("Local image-list execution cancelled") from exc
finally:
with self._state_lock:
self._loop = self._task = None
_EXECUTION_LOCK.release()
def _check_stopped(self):
import comfy.model_management
if self._stop.is_set() or comfy.model_management.processing_interrupted():
raise ImageListCancelledError("Local image-list execution interrupted")
async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
import execution
import torch
import comfy.model_management
import comfy.utils
from comfy_execution import progress
from comfy_execution.graph_utils import GraphBuilder
self._check_stopped()
with self._state_lock:
self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
dynamic_prompt = execution.DynamicPrompt(prompt)
private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
parent_progress = progress.global_progress_registry
parent_hook = comfy.utils.PROGRESS_BAR_HOOK
prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
progress.global_progress_registry = private_progress
def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
self._check_stopped()
# Never route a child progress/preview event through the parent's
# client or reset its handlers. Keep the private progress accurate.
if node_id is not None:
private_progress.update_progress(node_id, value, total, preview)
comfy.utils.PROGRESS_BAR_HOOK = progress_hook
try:
with torch.inference_mode():
if self._caches is None:
outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
objects = execution.HierarchicalCache(execution.CacheKeySetID)
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
caches = self._caches
changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
for cache in caches.all:
await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
cache.clean_unused()
execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
execution_list.add_node(end_id)
server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
assets = execution.NoAssets(execution.args)
extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
while not execution_list.is_empty():
self._check_stopped()
node_id, error, exc = await execution_list.stage_node_execution()
if error is not None:
raise RuntimeError(f"Local image-list graph failed: {error}") from exc
self._check_stopped()
result, error, exc = await execution.execute(
server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
execution_list, pending_subgraphs, pending_async, ui, assets)
if result == execution.ExecutionResult.FAILURE:
if isinstance(exc, (comfy.model_management.InterruptProcessingException,
ImageListCancelledError)):
raise ImageListCancelledError("Local image-list execution interrupted") from exc
raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
if result == execution.ExecutionResult.PENDING:
execution_list.unstage_node_execution()
else:
execution_list.complete_node_execution()
self._check_stopped()
finally:
# Core execute creates async-node and unblock tasks. This is a private
# event loop, so every other task is owned here. Drain/cancel them
# before restoring globals, including when the graph failed early.
with self._state_lock:
self._task = None # Stop requests must not interrupt this drain.
try:
while True:
pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
if not pending:
break
for task in pending:
task.cancel()
drain = asyncio.gather(*pending, return_exceptions=True)
while not drain.done():
try:
await asyncio.shield(drain)
except asyncio.CancelledError:
continue
finally:
progress.global_progress_registry = parent_progress
comfy.utils.PROGRESS_BAR_HOOK = parent_hook
for key, value in parent_prefix.items():
setattr(GraphBuilder, key, value)
File diff suppressed because it is too large Load Diff
+413
View File
@@ -0,0 +1,413 @@
"""Image-list items on explicitly selected comfy-media-v1 pools.
Only graph/media adaptation lives here. The shared manager owns admission,
provider submission, durable reservations, cancellation and atomic decoding.
A provider terminal state is not proof that an arbitrary endpoint stopped its
local Comfy prompt: the selected profile must honour its supervision contract.
"""
from __future__ import annotations
import asyncio
import copy
import hashlib
import uuid
from . import runtime
from ..api.orchestration.runpod_client import serialize_request
from ..api.orchestration.serverless_media_prompt import compile_media_prompt
from ..api.orchestration.serverless_runtime import PER_WORKER_JOB_LIMIT, get_manager
from ..api.orchestration.serverless_typed_inputs import prepare_media_inputs, _artifact_id
from ..serverless_bridge import codec, contract
from ..utils.serverless_config import validate_serverless_worker
def _identity(value):
contract.identifier(value)
if any(part in ("", ".", "..") for part in value.split("/")):
raise ValueError("Unsafe image-list media identifier")
return value
def resolve_serverless_workers(config, requested_ids):
"""Return pool configurations, never jobs_per_run-expanded participants.
None means enabled defaults; an explicit selection includes disabled pools.
Direct worker IDs may coexist in that selection and are left to their own
resolver. Unknown, duplicate configured and malformed identities fail closed.
"""
workers = (config or {}).get("workers", [])
indexed = {}
for worker in workers:
wid = worker.get("id")
if not isinstance(wid, str) or not wid or wid != wid.strip():
raise ValueError("Invalid configured image-list worker ID")
if wid in indexed:
raise ValueError("Duplicate configured image-list worker ID")
indexed[wid] = worker
requested = None
if requested_ids is not None:
if not isinstance(requested_ids, (list, tuple, set)):
raise ValueError("Image-list worker selection must be a list of exact IDs")
if any(not isinstance(wid, str) or not wid or wid != wid.strip() for wid in requested_ids):
raise ValueError("Invalid selected image-list worker ID")
requested = set(requested_ids)
if requested - indexed.keys():
raise ValueError("Unknown explicitly selected image-list worker ID(s)")
result = []
for wid, worker in indexed.items():
if requested is None:
if worker.get("enabled", False) is not True:
continue
elif wid not in requested:
continue
if str(worker.get("type") or "").strip().lower() != "serverless":
continue
clean = validate_serverless_worker(worker)
_identity(wid)
if clean.get("media_profile") != contract.PROFILE:
raise ValueError(f"Image-list serverless pool {wid!r} requires comfy-media-v1")
# Pool IDs, not virtual legacy participant IDs, own capacity and secrets.
if "pool_id" in clean and clean["pool_id"] != wid:
raise ValueError("Image-list selection requires a pool, not an expanded participant")
clean["pool_id"] = wid
result.append(clean)
return result
def _link(value, prompt):
if not isinstance(value, (list, tuple)):
return None
if (len(value) != 2 or not isinstance(value[0], str)
or type(value[1]) is not int or value[1] < 0):
raise ValueError("Invalid image-list graph link")
_identity(value[0])
if value[0] not in prompt:
raise ValueError("Image-list graph link source is missing")
return value[0], value[1]
def _new_node_id(graph, label):
base = "image-list:input:" + label
candidate, suffix = base, 0
while candidate in graph:
suffix += 1
candidate = f"{base}:{suffix}"
return candidate
def _merge_artifacts(*groups):
indexed = {}
for group in groups:
for artifact in group:
contract.validate_artifact(artifact)
aid = artifact["id"]
if aid in indexed and indexed[aid] != artifact:
raise ValueError("Conflicting image-list input artifact identity")
indexed[aid] = artifact
if (len(indexed) > contract.MAX_ARTIFACTS or
sum(a["byte_length"] for a in indexed.values()) > contract.MAX_TOTAL_BYTES):
raise ValueError("Image-list media inputs exceed artifact limits")
return [indexed[aid] for aid in sorted(indexed)]
async def prepare_serverless_item(prompt, region, item, worker, *, request_id, job_id):
"""Compile one immutable Start/End region to an endpoint-native media job.
Tensor Start outputs use exact codec artifacts and references to existing
consumer inputs. PrimitiveString placeholders keep exporter sockets linked
until the endpoint handler injects MediaImportV1 and prunes unused ancestors.
Scalar values remain typed literals, except exported text stays linked.
"""
_identity(request_id)
_identity(job_id)
worker = validate_serverless_worker(worker)
_identity(worker["id"])
if worker.get("media_profile") != contract.PROFILE:
raise ValueError("Image-list serverless execution requires comfy-media-v1")
if not isinstance(prompt, dict) or not prompt:
raise ValueError("Image-list prompt must be an API graph")
# Check original tokens before looking them up or pruning any nodes. Do not
# let int/bool IDs or slots become valid through string/int normalization.
for node_id, node in prompt.items():
_identity(node_id)
if not isinstance(node, dict) or not isinstance(node.get("inputs", {}), dict):
raise ValueError("Invalid image-list prompt node")
for name, value in node.get("inputs", {}).items():
_identity(name)
_link(value, prompt)
start, end = region["start_id"], region["end_id"]
_identity(start)
_identity(end)
if (prompt[start].get("class_type") != runtime.START_CLASS or
prompt[end].get("class_type") != runtime.END_CLASS):
raise ValueError("Image-list region has mismatched Start/End identities")
if _link(prompt[end].get("inputs", {}).get("context"), prompt) != (start, runtime.CONTEXT_OUT):
raise ValueError("Image-list End context must reference its exact Start")
ids = region["node_ids"]
for node_id in ids:
_identity(node_id)
if not {start, end} <= set(ids):
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")
image = item.get("image")
if (not isinstance(image, runtime._torch().Tensor) or image.ndim != 4
or image.shape[0] != 1 or image.shape[-1] not in (3, 4)):
raise ValueError("Image-list item must contain exactly one RGB/RGBA image frame")
values = item.get("values", {})
if not isinstance(values, dict) or any(name not in runtime.VALUE_NAMES for name in values):
raise ValueError("Invalid image-list value_N slots")
materialized, artifacts = {}, {}
placeholders = {}
def replacement(slot):
if slot in materialized:
return copy.deepcopy(materialized[slot])
if slot == runtime.IMAGE_OUT:
value, kind = image, "IMAGE"
elif slot == runtime.MASK_OUT:
value, kind = item.get("mask"), "MASK"
elif slot == runtime.INDEX_OUT:
value, kind = item["index"], "INT"
elif runtime.VALUE_OUT_BASE <= slot < runtime.VALUE_OUT_BASE + runtime.MAX_VALUE_SLOTS:
name = runtime.VALUE_NAMES[slot - runtime.VALUE_OUT_BASE]
value = values.get(name)
kind = runtime.value_kind(value, name)
if kind == "TENSOR":
if value.ndim == 4:
kind = "IMAGE"
elif value.ndim == 3:
kind = "MASK"
else:
raise ValueError(f"{name} serverless tensor requires IMAGE [B,H,W,C] or MASK [B,H,W]")
else:
raise ValueError("Unsupported Start context/control output in endpoint computation")
if value is None:
raise ValueError("Connected Start output is empty; None cannot be submitted")
if kind in ("IMAGE", "MASK"):
if kind == "MASK" and slot == runtime.MASK_OUT:
if value.ndim != 3 or value.shape[0] != 1 or value.shape[1:3] != image.shape[1:3]:
raise ValueError("Image-list item mask must be paired with its image grid")
if value.numel() * value.element_size() > contract.MAX_ARTIFACT_BYTES:
raise ValueError("Image-list tensor exceeds media artifact byte limit")
descriptor, payload = codec.encode_payload(kind, value.detach().to("cpu"),
max_bytes=contract.MAX_ARTIFACT_BYTES)
aid = _artifact_id(kind, descriptor, payload)
if aid not in artifacts and (len(artifacts) >= contract.MAX_ARTIFACTS or
len(payload) + sum(a["byte_length"] for a in artifacts.values()) > contract.MAX_TOTAL_BYTES):
raise ValueError("Image-list media inputs exceed artifact limits")
artifacts[aid] = contract.artifact_from_payload(aid, kind, descriptor, payload)
placeholder = _new_node_id(graph, str(slot))
graph[placeholder] = {"class_type": "PrimitiveString", "inputs": {"value": ""}}
placeholders[placeholder] = aid
output = [placeholder, 0]
else:
# value_kind rejects objects/None/nonfinite values and bool remains bool.
runtime.value_kind(value, "Start output")
output = value
materialized[slot] = output
return copy.deepcopy(output)
for node_id in list(graph):
inputs = graph[node_id].setdefault("inputs", {})
for name, value in list(inputs.items()):
link = _link(value, prompt)
if link and link[0] == start:
inputs[name] = replacement(link[1])
elif link:
inputs[name] = list(link)
# Start text/index/generic scalar substitutions may have removed End's link.
text = graph[end]["inputs"].get("texts")
if "texts" in graph[end]["inputs"] and not isinstance(text, list):
if not isinstance(text, str):
raise ValueError("Image-list End texts must be STRING")
nid = _new_node_id(graph, "text")
graph[nid] = {"class_type": "PrimitiveString", "inputs": {"value": text}}
graph[end]["inputs"]["texts"] = [nid, 0]
workflow, manifest = compile_media_prompt(graph, request_id)
refs = []
for nid, node in workflow.items():
for name, value in node["inputs"].items():
if isinstance(value, list) and value[0] in placeholders:
refs.append({"node_id": nid, "input_name": name, "artifact_id": placeholders[value[0]]})
# Only actual surviving consumer slots count as uses of an input artifact.
used = {ref["artifact_id"] for ref in refs}
images, fixed_artifacts, fixed_refs = await prepare_media_inputs(workflow)
merged = _merge_artifacts([artifacts[aid] for aid in used], fixed_artifacts)
refs.extend(fixed_refs)
media = {"version": contract.VERSION, "profile": contract.PROFILE,
"request_id": request_id, "artifacts": merged,
"inputs": sorted(refs, key=lambda ref: (ref["node_id"], ref["input_name"])),
"outputs": manifest}
# Includes handler identity/link/ref validation and aggregate + exact JSON caps.
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
return {"worker": worker, "workflow": workflow, "images": images,
"manifest": manifest, "media": media}
class ServerlessSlots:
"""Nonblocking slot adapter; completion means the manager task has drained."""
def __init__(self, state, region, workers, manager, event):
self.state, self.region, self.manager, self.event = state, region, manager, event
self.workers = {worker["id"]: worker for worker in workers}
# The manager remains the atomic authority at start(). Bound this local
# scheduler too, including other callers' currently pending tasks.
available = max(0, 64 - sum(not task.done() for task in manager.tasks.values())) if manager else 0
self.slots = [("serverless", worker["id"], index)
for index in range(PER_WORKER_JOB_LIMIT) for worker in workers][:available]
if workers and not self.slots:
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
digest = hashlib.sha256(state.context_id.encode()).hexdigest()
self._scope = f"image-list:{digest}:{uuid.uuid4().hex}"
self._active = {}
self._queues = {}
self._executions = set()
self._assigned = set()
self._submissions = set()
self._closed = False
self._close_task = None
async def dispatch(self, slot, index):
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
if slot not in self.slots or slot in self._active:
raise RuntimeError("Unknown or occupied serverless image-list slot")
if type(index) is not int or not 0 <= index < len(self.state.items) or index in self._assigned:
raise ValueError("Invalid or duplicate serverless image-list index")
if self.state.items[index].get("index") != index:
raise ValueError("Serverless image-list item index mismatch")
self._assigned.add(index)
identity = f"{self._scope}:{index}"
handle = {"index": index, "execution_id": identity, "job_id": identity,
"worker_id": slot[1], "task": None}
self._active[slot] = handle
# Own execution identity BEFORE start, even if its caller loses the handle.
self._executions.add(identity)
task = asyncio.create_task(self._submit(handle))
self._submissions.add(task)
try:
await asyncio.shield(task)
return handle
except BaseException:
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
self._remove_queue(handle["job_id"])
await self.manager.cancel(execution_id=identity)
raise
finally:
self._submissions.discard(task)
async def _submit(self, handle):
job = await prepare_serverless_item(
self.state.original_prompt, self.region, self.state.items[handle["index"]],
self.workers[handle["worker_id"]], request_id=handle["execution_id"], job_id=handle["job_id"])
if self._closed:
raise RuntimeError("Image-list serverless controller closed during preparation")
self.manager.ensure_capacity(1)
queue = asyncio.Queue()
async with self.manager.server.distributed_jobs_lock:
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
if handle["job_id"] in self.manager.server.distributed_pending_jobs:
raise RuntimeError("Duplicate serverless image-list queue identity")
self._queues[handle["job_id"]] = queue
self.manager.server.distributed_pending_jobs[handle["job_id"]] = queue
handle["sockets"] = {binding["socket"] for binding in job["manifest"]}
record_ids = await self.manager.start(handle["execution_id"], self.state.prompt_id, [job])
if len(record_ids) != 1 or record_ids[0] not in self.manager.tasks:
raise RuntimeError("Missing serverless image-list lifecycle task")
handle["task"] = self.manager.tasks[record_ids[0]]
handle["task"].add_done_callback(lambda _: self.event.set())
def _remove_queue(self, job_id):
queue = self._queues.pop(job_id, None)
# Same event-loop turn, without an await: never delete a replacement queue.
if queue is not None and self.manager.server.distributed_pending_jobs.get(job_id) is queue:
del self.manager.server.distributed_pending_jobs[job_id]
def poll(self, inflight):
if self._closed:
raise RuntimeError("Image-list serverless controller is closed")
for slot, handle in inflight.items():
if slot not in self.slots:
continue # Direct/master slots belong to the surrounding scheduler.
if self._active.get(slot) is not handle:
raise RuntimeError("Mismatched or duplicate serverless image-list completion handle")
task = handle["task"]
if task is None or not task.done():
continue
if task.cancelled():
raise RuntimeError("Serverless image-list task was cancelled")
if task.exception() is not None:
raise RuntimeError("Serverless image-list lifecycle task failed") from task.exception()
queue = self._queues.get(handle["job_id"])
if queue is None or self.manager.server.distributed_pending_jobs.get(handle["job_id"]) is not queue:
raise RuntimeError("Missing serverless image-list result queue")
if queue.qsize() != 1:
raise RuntimeError("Missing or duplicate serverless image-list terminal result")
message = queue.get_nowait()
if (not isinstance(message, dict) or message.get("worker_id") != handle["worker_id"]
or message.get("is_last") is not True):
raise RuntimeError("Mismatched serverless image-list result identity")
if "serverless_error" in message:
raise RuntimeError(f"Serverless image-list item {handle['index']} failed: {message['serverless_error']}")
media = message.get("media")
if not isinstance(media, dict) or set(media) != handle["sockets"]:
raise RuntimeError("Missing or unexpected serverless image-list result sockets")
if any(not isinstance(values, list) or len(values) != 1 or values[0] is None
for values in media.values()):
raise RuntimeError("Serverless image-list result requires exactly one value per socket")
result = runtime.normalize_worker_result(media["images"], media.get("masks"), media.get("texts"))
if result["mask"] is not None and result["mask"].shape[1:3] != result["image"].shape[1:3]:
raise ValueError("Serverless image-list result mask is not paired with the image grid")
self._remove_queue(handle["job_id"])
del self._active[slot]
return slot, handle["index"], result
return None
async def _close(self):
# Stop preparation/start before cancellation discovery, then retire exact
# destinations before manager callbacks can publish cancellation errors.
submissions = tuple(self._submissions)
for task in submissions:
if not task.done():
task.cancel()
await asyncio.gather(*submissions, return_exceptions=True)
for job_id in list(self._queues):
self._remove_queue(job_id)
await asyncio.gather(*(self.manager.cancel(execution_id=eid) for eid in self._executions))
self._active.clear()
async def close(self):
if self._close_task is None:
self._closed = True
self._close_task = asyncio.create_task(self._close())
cancelled = False
while not self._close_task.done():
try:
await asyncio.shield(self._close_task)
except asyncio.CancelledError:
cancelled = True
self._close_task.result()
if cancelled:
raise asyncio.CancelledError
async def create_serverless_slots(state, region, config, event):
workers = resolve_serverless_workers(config, state.worker_ids)
return ServerlessSlots(state, region, workers, get_manager() if workers else None, event)
+1 -31
View File
@@ -1,31 +1 @@
from .utilities import (
DistributedSeed,
DistributedModelName,
DistributedValue,
ImageBatchDivider,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from .collector import DistributedCollectorNode
NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"DistributedValue": DistributedValue,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed",
"DistributedModelName": "Distributed Model Name",
"DistributedValue": "Distributed Value",
"ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider",
"DistributedEmptyImage": "Distributed Empty Image",
}
"""Private execution helpers; public registration lives in nodes.v3."""
+457 -150
View File
@@ -12,23 +12,29 @@ from comfy.utils import ProgressBar
from ..utils.logging import debug_log, log
from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
from ..utils.constants import HEARTBEAT_INTERVAL
from ..utils.constants import HEARTBEAT_INTERVAL, TRANSFER_MAX_CHUNK_BYTES, TRANSFER_MAX_REQUEST_BYTES
from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
from ..utils.network import build_worker_url, get_client_session, probe_worker
from ..utils.audio_payload import encode_audio_payload
from ..utils.async_helpers import run_async_in_server_loop
from ..utils import transfer_codec
prompt_server = _server.PromptServer.instance
def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
"""Materialize every transfer chunk (runs in an executor)."""
return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
class DistributedCollectorNode:
INPUT_IS_LIST = True
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"load_balance": (
"BOOLEAN",
{
@@ -37,7 +43,27 @@ class DistributedCollectorNode:
},
),
},
"optional": { "audio": ("AUDIO",) },
"optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
"transfer_format": (
list(transfer_codec.TRANSFER_FORMATS),
{
"default": transfer_codec.DEFAULT_TRANSFER_FORMAT,
"tooltip": transfer_codec.TRANSFER_FORMAT_TOOLTIP,
},
),
"transfer_quality": (
"INT",
{
"default": transfer_codec.DEFAULT_TRANSFER_QUALITY,
"min": 1,
"max": 100,
"step": 1,
"tooltip": transfer_codec.TRANSFER_QUALITY_TOOLTIP,
},
),
},
"hidden": {
"multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}),
@@ -55,7 +81,76 @@ class DistributedCollectorNode:
FUNCTION = "run"
CATEGORY = "image"
def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
@staticmethod
def _unwrap_list_input(value):
"""Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
if isinstance(value, (list, tuple)) and len(value) == 1:
return value[0]
return value
def _normalize_images_input(self, images):
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
if isinstance(images, (list, tuple)):
if not images:
raise ValueError("Collector received an empty image list")
if not all(isinstance(image, torch.Tensor) for image in images):
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
if len(images) == 1:
return ensure_contiguous(images[0])
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
return ensure_contiguous(images)
def _normalize_audio_input(self, audio):
"""Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
if not isinstance(audio, (list, tuple)):
return audio
audio_items = [item for item in audio if item is not None]
if not audio_items:
return None
if len(audio_items) == 1:
return audio_items[0]
waveforms = []
sample_rate = 44100
for item in audio_items:
if not isinstance(item, dict):
raise TypeError("Collector expected AUDIO list items to be dictionaries")
waveform = item.get("waveform")
if waveform is None or waveform.numel() == 0:
continue
waveforms.append(waveform)
if sample_rate == 44100:
sample_rate = item.get("sample_rate", 44100)
if not waveforms:
return None
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
if images is not None:
images = self._normalize_images_input(images)
audio = self._normalize_audio_input(audio)
load_balance = self._unwrap_list_input(load_balance)
multi_job_id = self._unwrap_list_input(multi_job_id)
is_worker = self._unwrap_list_input(is_worker)
master_url = self._unwrap_list_input(master_url)
enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
worker_batch_size = self._unwrap_list_input(worker_batch_size)
worker_id = self._unwrap_list_input(worker_id)
pass_through = self._unwrap_list_input(pass_through)
delegate_only = self._unwrap_list_input(delegate_only)
transfer_format = self._unwrap_list_input(transfer_format)
transfer_quality = self._unwrap_list_input(transfer_quality)
remote_only_master = (
bool(multi_job_id)
and not is_worker
and (delegate_only or is_master_delegate_only())
)
if images is None and audio is None and not remote_only_master:
raise ValueError("DistributedCollector requires at least one image or audio input")
# Create empty audio if not provided
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -77,46 +172,218 @@ class DistributedCollectorNode:
worker_batch_size,
worker_id,
delegate_only,
transfer_format,
transfer_quality,
)
)
return result
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send image batch to master via canonical JSON envelopes."""
batch_size = image_batch.shape[0]
if batch_size == 0:
return
@staticmethod
def _get_serverless_manager():
"""Return the serverless manager only for collector jobs with serverless slots."""
manager = getattr(prompt_server, "distributed_serverless", None)
if manager is not None:
return manager
from ..api.orchestration.serverless_runtime import get_manager
return get_manager()
async def send_batch_to_master(
self,
image_batch,
audio,
multi_job_id,
master_url,
worker_id,
transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT,
transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY,
):
"""Send an image batch, optionally with audio, or an audio-only completion."""
batch_size = 0 if image_batch is None else image_batch.shape[0]
fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
if fmt != transfer_codec.LEGACY_FORMAT:
sent = await self._send_compressed_batch(
image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
)
if sent:
return
log(
"Worker - Master did not accept compressed transfer "
"(older ComfyUI-Distributed on the master?). Falling back to legacy_png."
)
await self._send_legacy_batch(image_batch, audio, multi_job_id, master_url, worker_id)
@staticmethod
def _resolve_transfer_format(transfer_format):
try:
fmt = transfer_codec.normalize_format(transfer_format)
except transfer_codec.TransferCodecError as exc:
log(f"Worker - {exc}; using legacy_png.")
return transfer_codec.LEGACY_FORMAT
if fmt not in transfer_codec.VIDEO_FORMATS or transfer_codec.video_codec_available(fmt):
return fmt
fallback = "h264" if fmt != "ffv1" and transfer_codec.video_codec_available("h264") else "png"
log(f"Worker - Encoder for '{fmt}' is not available in this PyAV build; using '{fallback}'.")
return fallback
@staticmethod
def _build_multipart(envelope, parts):
form = aiohttp.FormData()
form.add_field("envelope", json.dumps(envelope), content_type="application/json")
for i, data in enumerate(parts):
form.add_field("media", data, filename=f"{i}.bin", content_type="application/octet-stream")
return form
# Room left in each request for the multipart envelope (ids, frame
# metadata, part headers). Audio is budgeted separately.
TRANSFER_ENVELOPE_RESERVE_BYTES = 256 * 1024
# Body of a 400 from a master without compressed-transfer support: it
# tries request.json() on the multipart body and fails.
_OLD_MASTER_REJECTION = "Invalid JSON payload"
async def _send_compressed_batch(self, image_batch, audio, multi_job_id, master_url, worker_id, fmt, quality):
"""Send the batch as compressed multipart chunks.
All chunks are encoded before anything is sent, so a frame that cannot
fit in one request fails the transfer without delivering a partial
batch. Every request stays within TRANSFER_MAX_REQUEST_BYTES.
Returns False (without having delivered anything) when the master
rejects the first chunk because it predates compressed transfer, so
the caller can retry with legacy_png. Any other rejection raises.
"""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
loop = asyncio.get_running_loop()
frames = await loop.run_in_executor(None, transfer_codec.tensor_batch_to_uint8, image_batch)
raw_bytes = int(frames.nbytes)
total = int(frames.shape[0])
media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
started = time.time()
# Encoding is CPU-bound (seconds for video); keep the server loop responsive.
chunks = await loop.run_in_executor(
None, _encode_all_chunks, frames, fmt, quality, min(TRANSFER_MAX_CHUNK_BYTES, media_limit), media_limit,
)
del frames
# Attach audio to the final chunk only if the request stays in bounds;
# otherwise finish with a separate audio-only completion.
audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
last_media = sum(len(p) for p in chunks[-1]["parts"])
audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
audio_separate = encoded_audio is not None and not audio_inline
sent_bytes = 0
for index, chunk in enumerate(chunks):
parts = chunk.pop("parts")
is_final_chunk = index == len(chunks) - 1
envelope = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
"is_last": bool(is_final_chunk and not audio_separate),
**chunk,
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
if is_final_chunk and audio_inline:
envelope["audio"] = encoded_audio
chunk_bytes = sum(len(p) for p in parts)
timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
data=self._build_multipart(envelope, parts),
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
if response.status >= 400:
body = await response.text()
if index == 0 and response.status == 400 and self._OLD_MASTER_REJECTION in body:
debug_log(f"Worker - Master rejected compressed chunk: {body[:300]}")
return False
raise RuntimeError(
f"Master rejected {fmt} chunk {index + 1}/{len(chunks)} "
f"({chunk_bytes} bytes) with HTTP {response.status}: {body[:300]}. "
"If this is a size limit, make COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES "
"on the worker smaller than the master's --max-upload-size."
)
except Exception as e:
log(f"Worker - Failed to send canonical image envelope to master: {e}")
log(f"Worker - Failed to send {fmt} chunk to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise # Re-raise to handle at caller level
raise
sent_bytes += chunk_bytes
requests = len(chunks)
if audio_separate:
await self._post_json_completion(session, url, {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
})
requests += 1
log(
f"Worker - Sent {total} image(s) as {fmt} in {requests} request(s): "
f"{sent_bytes / 1e6:.2f} MB (raw {raw_bytes / 1e6:.2f} MB, "
f"{transfer_codec.describe_ratio(sent_bytes, raw_bytes)} smaller) in {time.time() - started:.1f}s"
)
return True
@staticmethod
async def _post_json_completion(session, url, payload):
timeout_seconds = 60 if "image" in payload else 600
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
except Exception as e:
media_type = "image/audio" if "image" in payload else "audio-only"
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
debug_log(f"Worker - Full error details: URL={url}")
raise
async def _send_legacy_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Original transfer: one JSON request per frame with a base64 PNG."""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
payloads = []
batch_size = 0 if image_batch is None else image_batch.shape[0]
if batch_size == 0:
if encoded_audio is None:
raise ValueError("Worker completion requires image or audio data")
payloads.append(
{
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
}
)
else:
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
payloads.append(payload)
for payload in payloads:
await self._post_json_completion(session, url, payload)
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output.
@@ -198,8 +465,8 @@ class DistributedCollectorNode:
images_on_cpu,
delegate_mode: bool,
fallback_images,
) -> torch.Tensor:
"""Assemble final tensor: master first, then workers in enabled order."""
):
"""Assemble final tensor, or return None when the job contains only audio."""
ordered_tensors = []
if not delegate_mode and images_on_cpu is not None:
for i in range(master_batch_size):
@@ -230,16 +497,18 @@ class DistributedCollectorNode:
if cpu_tensors:
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
elif fallback_images is not None:
if fallback_images is not None:
return ensure_contiguous(fallback_images)
else:
raise ValueError("No image data collected from master or workers")
return None
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
if is_worker:
# Worker mode: send images and audio to master in a single batch
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
image_count = 0 if images is None else images.shape[0]
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
await self.send_batch_to_master(
images, audio, multi_job_id, master_url, worker_id, transfer_format, transfer_quality
)
return (images, audio if audio is not None else self.EMPTY_AUDIO)
else:
delegate_mode = delegate_only or is_master_delegate_only()
@@ -254,6 +523,17 @@ class DistributedCollectorNode:
seen_enabled.add(worker_id_str)
enabled_workers.append(worker_id_str)
expected_workers = set(enabled_workers)
serverless_workers = {
worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
}
serverless_manager = None
def get_serverless_manager():
nonlocal serverless_manager
if serverless_workers and serverless_manager is None:
serverless_manager = self._get_serverless_manager()
return serverless_manager
num_workers = len(expected_workers)
if num_workers == 0:
return (images, audio if audio is not None else self.EMPTY_AUDIO)
@@ -273,14 +553,15 @@ class DistributedCollectorNode:
master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else:
images_on_cpu = images.cpu()
master_batch_size = images.shape[0]
if images is None:
images_on_cpu = None
master_batch_size = 0
else:
images_on_cpu = ensure_contiguous(images.cpu())
master_batch_size = images.shape[0]
master_audio = audio # Keep master's audio for later
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
# Ensure master images are contiguous
images_on_cpu = ensure_contiguous(images_on_cpu)
# Initialize storage for collected images and audio
worker_images = {} # Dict to store images by worker_id and index
@@ -319,151 +600,177 @@ class DistributedCollectorNode:
workers_done.add(done_worker_id)
p.update(1) # +1 per completed expected worker
def consume_result(result):
nonlocal collected_count
if "serverless_error" in result:
error = result.get("serverless_error")
message = str(error).strip() if error is not None else "unspecified serverless failure"
participant = result.get("worker_id", "unknown")
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
result_worker_id = str(result["worker_id"])
is_last = result.get("is_last", False)
collected_count += self._store_worker_result(worker_images, result)
debug_log(
f"Master - Got canonical result from worker {result_worker_id}, "
f"image {result.get('image_index', 0)}, is_last={is_last}"
)
result_audio = result.get("audio")
if result_audio is not None:
worker_audio[result_worker_id] = result_audio
debug_log(f"Master - Got audio from worker {result_worker_id}")
if is_last:
mark_worker_done(result_worker_id)
try:
while len(workers_done) < num_workers:
# Check for user interruption to abort collection promptly
# Check for user interruption to abort collection promptly.
comfy.model_management.throw_exception_if_processing_interrupted()
try:
# Get the queue again each time to ensure we have the right reference
async with prompt_server.distributed_jobs_lock:
q = prompt_server.distributed_pending_jobs[multi_job_id]
current_size = q.qsize()
result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
count = self._store_worker_result(worker_images, result)
collected_count += count
debug_log(
f"Master - Got canonical result from worker {worker_id}, "
f"image {result.get('image_index', 0)}, is_last={is_last}"
)
consume_result(result)
# Collect audio data if present
result_audio = result.get('audio')
if result_audio is not None:
worker_audio[worker_id] = result_audio
debug_log(f"Master - Got audio from worker {worker_id}")
# Record activity and refresh timeout baseline
# Real results refresh the regular-worker timeout baseline.
last_activity = time.time()
base_timeout = float(get_worker_timeout_seconds())
if is_last:
mark_worker_done(worker_id)
except asyncio.TimeoutError:
# If we still have time, continue polling; otherwise handle timeout
if (time.time() - last_activity) < base_timeout:
comfy.model_management.throw_exception_if_processing_interrupted()
continue
# Re-check for user interruption after timeout expiry
comfy.model_management.throw_exception_if_processing_interrupted()
missing_workers = set(str(w) for w in enabled_workers) - workers_done
missing_workers = expected_workers - workers_done
# A queue delivery can race the timed wait. Process it before
# declaring an owned provider slot stopped without a result.
async with prompt_server.distributed_jobs_lock:
current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
queue_has_results = current_q is not None and not current_q.empty()
if queue_has_results:
continue
missing_serverless = missing_workers & serverless_workers
if missing_serverless:
manager = get_serverless_manager()
if manager is None:
raise RuntimeError("Serverless manager is unavailable for pending collector slots")
pending_serverless = []
for participant_id in sorted(missing_serverless):
if manager.is_pending(multi_job_id, participant_id):
pending_serverless.append(participant_id)
elif manager.owns(multi_job_id, participant_id):
raise RuntimeError(
f"Serverless participant {participant_id} stopped without delivering a result"
)
# Provider deadlines are managed by the serverless runtime;
# do not apply the regular worker heartbeat timeout to it.
if pending_serverless:
continue
missing_regular = missing_workers - serverless_workers
elapsed = time.time() - last_activity
for missing_worker_id in sorted(missing_workers):
for missing_worker_id in sorted(missing_regular):
log(
"Master - Heartbeat timeout: "
f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
)
log(
f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
f"(elapsed={elapsed:.1f}s)"
)
if missing_regular:
log(
f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
f"(elapsed={elapsed:.1f}s)"
)
# Probe missing workers' /prompt endpoints to check if they are actively processing
# Probe only regular worker /prompt endpoints.
any_busy = False
try:
cfg = load_config()
cfg_workers = cfg.get('workers', [])
for wid in list(missing_workers):
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
if not wrec:
debug_log(f"Collector probe: worker {wid} not found in config")
continue
worker_url = build_worker_url(wrec)
try:
payload = await probe_worker(worker_url, timeout=2.0)
queue_remaining = None
if payload is not None:
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
debug_log(
"Collector probe: worker "
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
)
if payload is not None and queue_remaining and queue_remaining > 0:
any_busy = True
log(
f"Master - Probe grace: worker {wid} appears busy "
f"(queue_remaining={queue_remaining}). Continuing to wait."
if missing_regular:
try:
cfg = load_config()
cfg_workers = cfg.get('workers', [])
for wid in sorted(missing_regular):
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
if not wrec:
debug_log(f"Collector probe: worker {wid} not found in config")
continue
worker_url = build_worker_url(wrec)
try:
payload = await probe_worker(worker_url, timeout=2.0)
queue_remaining = None
if payload is not None:
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
debug_log(
"Collector probe: worker "
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
)
break
except Exception as e:
debug_log(f"Collector probe failed for worker {wid}: {e}")
except Exception as e:
debug_log(f"Collector probe setup error: {e}")
if payload is not None and queue_remaining and queue_remaining > 0:
any_busy = True
log(
f"Master - Probe grace: worker {wid} appears busy "
f"(queue_remaining={queue_remaining}). Continuing to wait."
)
break
except Exception as e:
debug_log(f"Collector probe failed for worker {wid}: {e}")
except Exception as e:
debug_log(f"Collector probe setup error: {e}")
if any_busy:
# Refresh last_activity and continue waiting
last_activity = time.time()
# Refresh base timeout in case the user changed it in UI
base_timeout = float(get_worker_timeout_seconds())
continue
# Check queue size again with lock
async with prompt_server.distributed_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_jobs:
final_q = prompt_server.distributed_pending_jobs[multi_job_id]
final_size = final_q.qsize()
# Try to drain any remaining items
remaining_items = []
final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
remaining_items = []
if final_q is not None:
while not final_q.empty():
try:
item = final_q.get_nowait()
remaining_items.append(item)
remaining_items.append(final_q.get_nowait())
except asyncio.QueueEmpty:
break
if remaining_items:
# Process them
for item in remaining_items:
worker_id = item['worker_id']
is_last = item.get('is_last', False)
collected_count += self._store_worker_result(worker_images, item)
if is_last:
mark_worker_done(worker_id)
else:
log(f"Master - Queue {multi_job_id} no longer exists!")
for item in remaining_items:
consume_result(item)
break
except comfy.model_management.InterruptProcessingException:
# Cleanup queue on interruption and re-raise to abort prompt cleanly
total_collected = sum(len(imgs) for imgs in worker_images.values())
# Remove the queue before combining so a late provider callback cannot
# resurrect results for a collector that has completed.
async with prompt_server.distributed_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_jobs:
del prompt_server.distributed_pending_jobs[multi_job_id]
raise
total_collected = sum(len(imgs) for imgs in worker_images.values())
# Clean up job queue
async with prompt_server.distributed_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_jobs:
del prompt_server.distributed_pending_jobs[multi_job_id]
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
try:
combined = self._reorder_and_combine_tensors(
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
)
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
# Combine audio from master and workers
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
return (combined, combined_audio)
except Exception as e:
log(f"Master - Error combining images: {e}")
# Return just the master images as fallback
return (images, audio if audio is not None else self.EMPTY_AUDIO)
serverless_results = set(worker_images) & serverless_workers
try:
combined = self._reorder_and_combine_tensors(
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
)
if combined is None:
debug_log(f"Master - Job {multi_job_id} complete with audio only")
else:
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined, combined_audio)
except Exception as e:
if serverless_results:
raise RuntimeError(
"Serverless result batch shape mismatch while combining participants "
f"{sorted(serverless_results)}: {e}"
) from e
log(f"Master - Error combining images: {e}")
# Preserve collected audio for regular workers when image assembly fails.
return (images, combined_audio)
except BaseException:
# Remove the queue first, preventing cancelled work from delivering into
# a collector that is already unwinding.
async with prompt_server.distributed_jobs_lock:
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
if serverless_workers:
try:
manager = get_serverless_manager()
if manager is not None:
await manager.cancel(collector_job_id=multi_job_id)
except Exception as cancel_error:
debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
raise
+6 -15
View File
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
"""
Distributed version of Ultimate SD Upscale (No Upscale).
Supports three processing modes:
Supports two currently selected processing modes:
1. Single GPU: No workers available, process everything locally
2. Static Mode: Small batches, distributes tiles across workers (flattened)
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
2. Distributed tile queue: Workers pull tiles from a shared queue
Features:
- Multi-mode batch handling for efficient video/image upscaling
- Tile-based batch handling for video/image upscaling
- Tiled VAE support for memory efficiency
- Dynamic load balancing for large batches
- Shared work queue so faster workers can process more tiles
- Backward compatible with single-image workflows
Environment Variables:
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
Threshold: dynamic_threshold input controls mode switch (default 8)
The hidden dynamic_threshold input is retained for workflow compatibility but
does not affect the current mode-selection policy.
"""
def __init__(self):
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
# Ensure initialization before registering routes
ensure_tile_jobs_initialized()
# Node registration
NODE_CLASS_MAPPINGS = {
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
}
+268
View File
@@ -0,0 +1,268 @@
"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
from __future__ import annotations
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,
VALUES_INPUT,
create_master_job,
normalize_image_list_inputs,
normalize_worker_result,
run_image_list_job,
encode_result,
encode_request_payload,
unwrap_single,
)
# Types the transport can carry exactly. Anything else (models, latents, ...)
# must be loaded inside the region or shared as a model dependency instead.
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
assert len(VALUE_NAMES) == MAX_VALUE_SLOTS
class DistributedImageListStart(io.ComfyNode):
"""Materialize a batch once, then expose ordered per-item typed values."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListStart",
display_name="Distributed Image List Start",
category="image/distributed",
description=("Starts one per-image processing region. Each image becomes one job. "
"mask and value_N inputs take 1 value (shared by every image) or exactly "
"one value per image, paired in order."),
inputs=[
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one job."),
io.Mask.Input("mask", optional=True,
tooltip="1 mask shared by every image, or one mask per image."),
io.Autogrow.Input(VALUES_INPUT, optional=True, template=io.Autogrow.TemplateNames(
io.MultiType.Input("value", types=_VALUE_TYPES,
tooltip=("Any image, mask, text, number or boolean. 1 value is shared by "
"every image; one value per image pairs them in order. "
"A batch stays one value.")),
names=list(VALUE_NAMES), min=0)),
],
outputs=[
io.Image.Output(display_name="image", is_output_list=True),
io.Mask.Output(display_name="mask", is_output_list=True),
io.Int.Output(display_name="item_index", is_output_list=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
# value_N outputs mirror value_N inputs. The web extension hides
# unused ones and shows the connected input's type.
*[io.AnyType.Output(id=name, display_name=name, is_output_list=True) for name in VALUE_NAMES],
],
is_input_list=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
)
@classmethod
def execute(cls, images=None, mask=None, values=None,
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
image_list_master_callback="", image_list_execution_id="",
image_list_workflow_meta=None, image_list_master_participates=None,
image_list_local_execution=False):
worker_item_token = unwrap_single(worker_item_token)
worker_context_id = unwrap_single(worker_context_id)
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
if worker_item_token:
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import consume_worker_item
item = consume_worker_item(worker_item_token)
else:
item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
context_id = worker_context_id
if not isinstance(context_id, str) or not context_id:
raise RuntimeError("Worker image-list prompt is missing its paired context")
item_values = item.get("values") or {}
return io.NodeOutput(
[item["image"]], [item.get("mask")], [item["index"]], [context_id],
*[[item_values.get(name)] for name in VALUE_NAMES],
)
items = normalize_image_list_inputs(images, mask, values)
prompt = getattr(cls.hidden, "prompt", None)
start_id = str(getattr(cls.hidden, "unique_id", ""))
if not prompt or not start_id:
raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
try:
worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
except (TypeError, ValueError) as exc:
raise RuntimeError("Invalid preflighted image-list worker list") from exc
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
raise RuntimeError("Invalid preflighted image-list worker list")
from comfy_execution.utils import get_executing_context
from ..image_list.lifecycle import has_safe_execution_targets
import server
executing = get_executing_context()
explicit_targets = has_safe_execution_targets(
getattr(server.PromptServer, "instance", None),
executing.prompt_id if executing is not None else None, prompt)
from ..image_list.runtime import validate_image_list_regions
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets)
region = next((row for row in regions if row["start_id"] == start_id), None)
if region is None:
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
master_participates = unwrap_single(image_list_master_participates)
if master_participates is None:
from ..utils.config import load_config
master_participates = not bool((load_config().get("settings", {}) or {}).get("master_delegate_only", False))
if not isinstance(master_participates, bool):
raise ValueError("image_list_master_participates must be a boolean")
context_id = create_master_job(
start_id, region["end_id"], prompt, items, worker_ids,
str(unwrap_single(image_list_master_callback or "")),
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
prompt_id=executing.prompt_id if executing is not None else None,
master_participates=master_participates,
explicit_targets=explicit_targets,
)
return io.NodeOutput(
[item["image"] for item in items],
[item.get("mask") for item in items],
[item["index"] for item in items],
[context_id for _ in items],
*[[item["values"].get(name) for item in items] for name in VALUE_NAMES],
)
async def _consume_worker_item(token):
# Item state is local to each worker process; source nodes consume it exactly once.
from ..image_list.runtime import consume_worker_item
return consume_worker_item(token)
class DistributedImageListEnd(io.ComfyNode):
"""Collect branch results in source order without evaluating the branch on master."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListEnd",
display_name="Distributed Image List End",
category="image/distributed",
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
inputs=[
io.Image.Input("images", lazy=True),
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),
io.Mask.Output(display_name="masks", is_output_list=True),
io.String.Output(display_name="texts", is_output_list=True),
],
is_input_list=True,
is_output_node=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id, io.Hidden.execution_list],
)
@classmethod
def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", **kwargs):
worker_result_token = unwrap_single(worker_result_token)
if worker_result_token:
# Decide which optional sockets are connected from the actual worker prompt;
# absent optional inputs must not be demanded from Comfy's lazy executor.
prompt = getattr(cls.hidden, "prompt", {}) or {}
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
if isinstance(node, dict)
and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
required = []
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
connected = name == "images" or name in end_inputs
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
if connected and missing:
required.append(name)
return required
# Master only needs Start's context. Never request branch outputs here.
if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
return ["context"]
return []
@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, 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:
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Worker Image List End is missing its paired context")
result = normalize_worker_result(images, masks, texts)
if unwrap_single(image_list_local_execution):
from ..image_list.runtime import set_master_result
set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
str(worker_result_token), result)
else:
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,
)
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Distributed Image List End requires the paired Start context wire")
context_id = context_values[0]
from ..image_list.runtime import ImageListCancelledError
try:
# The runtime owns the configurable deadline and bounded remote cleanup.
execution_list = getattr(cls.hidden, "execution_list", None)
# Local model outputs belong to this still-running outer prompt.
# Add them to its existing tracker; never start/end another tracker
# that could unpin a model still needed by the parent.
model_callback = getattr(execution_list, "output_link_callback", None)
result_lists = run_async_in_server_loop(run_image_list_job(context_id, model_callback=model_callback))
except ImageListCancelledError:
import comfy.model_management
raise comfy.model_management.InterruptProcessingException()
images_out, masks_out, texts_out = result_lists
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
return io.NodeOutput(images_out, masks_out, texts_out)
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
if not callback_url.startswith(("http://", "https://")):
raise RuntimeError("Worker image-list callback URL is missing or invalid")
import aiohttp
from ..utils.network import get_client_session
session = await get_client_session()
payload = {"context_id": context_id, "token": item_token,
"item_index": item_index, "result": result}
async with session.post(
f"{callback_url.rstrip('/')}/distributed/image-list/result",
data=encode_request_payload(payload),
headers={"Content-Type": "application/json"},
timeout=aiohttp.ClientTimeout(total=60),
) as response:
if response.status >= 400:
body = (await response.text())[:500]
raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")
+229
View File
@@ -0,0 +1,229 @@
"""Typed, lossless list collection for serverless media participants.
The master compiler supplies expected_media before pruning delegate inputs.
The provider manager owns remote deadlines and delivers one atomic native-media
message per participant. This node never serializes, concatenates or resamples.
"""
import asyncio
import json
import time
import comfy.model_management
from comfy_api.v0_0_2 import io
import server
from ..serverless_bridge.contract import MAX_ITEMS, SOCKET_KINDS, identifier
from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
from ..utils.async_helpers import run_async_in_server_loop
POLL_SECONDS = 0.1
START_GRACE_SECONDS = 2.0
def _control(value, name, kind):
# Only execution controls are scalars wrapped by ComfyUI's list mapper.
if isinstance(value, list):
if len(value) != 1:
raise ValueError(f'{name} must be a singleton control')
value = value[0]
if type(value) is not kind:
raise ValueError(f'Invalid {name} control')
return value
def _names(value, name):
try:
result = json.loads(_control(value, name, str))
except (ValueError, TypeError):
raise ValueError(f'{name} must be a JSON string list') from None
if (not isinstance(result, list) or any(not isinstance(item, str) for item in result)
or len(result) != len(set(result))):
raise ValueError(f'{name} must contain unique strings')
return result
def _validate_media(media, expected):
if not isinstance(media, dict) or set(media) != set(expected):
raise ValueError('Media sockets do not match expected_media')
for socket, items in media.items():
if (socket not in SOCKET_KINDS or not isinstance(items, list)
or not 1 <= len(items) <= MAX_ITEMS or any(item is None for item in items)):
raise ValueError('Connected media sockets require nonempty bounded item lists')
def _prepare(media, controls):
worker = _control(controls['is_worker'], 'is_worker', bool)
delegate = _control(controls['delegate_only'], 'delegate_only', bool)
passthrough = _control(controls['pass_through'], 'pass_through', bool)
_control(controls['load_balance'], 'load_balance', bool) # Routing is master-owned.
if worker:
raise ValueError('DistributedMediaCollector does not support direct-worker execution')
participants = _names(controls['enabled_worker_ids'], 'enabled_worker_ids')
for participant in participants:
identifier(participant)
if not participant.startswith('serverless:') or not participant[len('serverless:'):]:
raise ValueError('DistributedMediaCollector requires serverless participant IDs')
expected = _names(controls['expected_media'], 'expected_media')
if any(socket not in SOCKET_KINDS for socket in expected):
raise ValueError('Unknown expected_media socket')
local = {socket: items for socket, items in media.items() if items is not None}
_validate_media(local, local)
expected = expected or list(local)
if not expected:
raise ValueError('At least one connected media socket is required')
if not delegate or passthrough:
_validate_media(local, expected)
if delegate and not passthrough and not participants:
raise ValueError('Delegate-only collection requires serverless participants')
return local, participants, expected, delegate, passthrough
def _outputs(media):
return tuple(media.get(socket, []) for socket in SOCKET_KINDS)
def _consume(message, participants, expected, received):
if not isinstance(message, dict):
raise RuntimeError('Invalid serverless media delivery')
participant = message.get('worker_id')
if not isinstance(participant, str) or participant not in participants:
raise RuntimeError('Unexpected serverless participant')
if participant in received:
raise RuntimeError('Duplicate serverless participant delivery')
if message.get('is_last') is not True:
raise RuntimeError('Serverless media delivery must be atomic and final')
if 'serverless_error' in message:
if (set(message) != {'worker_id', 'serverless_error', 'is_last'}
or not isinstance(message['serverless_error'], str) or not message['serverless_error']):
raise RuntimeError('Invalid serverless error delivery')
# Do not echo provider-controlled text, credentials or signed URLs.
raise RuntimeError(f'Serverless participant {participant} failed to produce media')
if set(message) != {'worker_id', 'media', 'is_last'}:
raise RuntimeError('Invalid serverless media delivery fields')
_validate_media(message['media'], expected)
received[participant] = message['media']
async def _collect(job, media, controls):
ps = server.PromptServer.instance
queue = None
try:
# Reuse the master's precreated queue, including early completions.
async with ps.distributed_jobs_lock:
queue = ps.distributed_pending_jobs.setdefault(job, asyncio.Queue())
local, participants, expected, delegate, passthrough = _prepare(media, controls)
if passthrough or not participants:
return _outputs(local)
received = {}
started = time.monotonic()
while len(received) < len(participants):
comfy.model_management.throw_exception_if_processing_interrupted()
async with ps.distributed_jobs_lock:
if ps.distributed_pending_jobs.get(job) is not queue:
raise RuntimeError('Serverless collector queue was removed or replaced')
try:
message = await asyncio.wait_for(queue.get(), timeout=POLL_SECONDS)
except asyncio.TimeoutError:
comfy.model_management.throw_exception_if_processing_interrupted()
# wait_for cancellation may race an atomic manager delivery.
if not queue.empty():
continue
manager = getattr(ps, 'distributed_serverless', None)
for participant in participants:
if participant in received:
continue
if manager is not None and manager.is_pending(job, participant):
continue # The manager enforces its remote absolute deadline.
if manager is not None and manager.owns(job, participant):
raise RuntimeError('Serverless participant stopped without delivering media')
# The master queues its prompt before manager.start journals slots.
if time.monotonic() - started >= START_GRACE_SECONDS:
raise RuntimeError('Serverless manager did not start the expected participant')
continue
_consume(message, participants, expected, received)
comfy.model_management.throw_exception_if_processing_interrupted()
# Drain and close atomically with the manager's delivery lock. A queued
# duplicate after the final expected message must not yield success.
async with ps.distributed_jobs_lock:
if ps.distributed_pending_jobs.get(job) is not queue:
raise RuntimeError('Serverless collector queue was removed or replaced')
while not queue.empty():
_consume(queue.get_nowait(), participants, expected, received)
del ps.distributed_pending_jobs[job]
ordered = [] if delegate else [local]
ordered.extend(received[participant] for participant in participants)
return tuple([item for source in ordered for item in source.get(socket, [])]
for socket in SOCKET_KINDS)
except BaseException:
# Remove only our queue before cancellation can trigger final deliveries.
async with ps.distributed_jobs_lock:
if queue is not None and ps.distributed_pending_jobs.get(job) is queue:
del ps.distributed_pending_jobs[job]
manager = getattr(ps, 'distributed_serverless', None)
if manager is not None:
try:
await manager.cancel(collector_job_id=job)
except BaseException:
pass # Preserve the original failure/interruption, not cleanup's.
raise
finally:
async with ps.distributed_jobs_lock:
if queue is not None and ps.distributed_pending_jobs.get(job) is queue:
del ps.distributed_pending_jobs[job]
async def _bridge_result(coro):
# The legacy sync bridge catches Exception, not asyncio.CancelledError.
# Carry BaseException as data so cancellation is re-raised on the caller.
try:
return await coro, None
except BaseException as error:
return None, error
class DistributedMediaCollector(io.ComfyNode):
@classmethod
def define_schema(cls):
types = (io.Image, io.Audio, io.Mask, io.String, io.Video)
return io.Schema(
node_id='DistributedMediaCollector', display_name='Distributed Media Collector',
category='utils',
inputs=[io.Boolean.Input('load_balance', default=False),
# Native non-widget inputs expose forceInput via extra_dict.
*[kind.Input(socket, optional=True,
**({'force_input': True} if socket == 'texts'
else {'extra_dict': {'forceInput': True}}))
for socket, kind in zip(SOCKET_KINDS, types)],
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=[kind.Output(display_name=socket, is_output_list=True)
for socket, kind in zip(SOCKET_KINDS, types)],
is_input_list=True, accept_all_inputs=True,
)
@classmethod
def execute(cls, images=None, audio=None, masks=None, texts=None, video=None,
load_balance=False, multi_job_id='', is_worker=False,
enabled_worker_ids='[]', delegate_only=False, pass_through=False,
expected_media='[]', transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
validate_options(_control(transfer_format, 'transfer_format', str),
_control(transfer_quality, 'transfer_quality', int))
media = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
controls = dict(load_balance=load_balance, is_worker=is_worker,
enabled_worker_ids=enabled_worker_ids, delegate_only=delegate_only,
pass_through=pass_through, expected_media=expected_media)
job = _control(multi_job_id, 'multi_job_id', str)
if not job:
local, participants, _, _, passthrough = _prepare(media, controls)
if participants and not passthrough:
raise ValueError('Serverless collection requires multi_job_id')
return io.NodeOutput(*_outputs(local))
identifier(job)
result, error = run_async_in_server_loop(_bridge_result(_collect(job, media, controls)))
if error is not None:
raise error
return io.NodeOutput(*result)
+2 -2
View File
@@ -269,7 +269,7 @@ class ImageBatchDivider:
class AudioBatchDivider:
"""Divides an audio waveform into multiple parts along the time/samples dimension."""
"""Divides an audio waveform into sequential segments along the time dimension."""
@classmethod
def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the audio into"
"tooltip": "Number of sequential time segments to create"
}),
}
}
+216
View File
@@ -0,0 +1,216 @@
"""Native V3 schemas with the existing execution algorithms kept intact.
Runtime objects are private, per-execution helpers, not registered V1 nodes.
This avoids sharing mutable instance state through V3's sanitized class clones.
"""
import comfy.samplers
from comfy_api.v0_0_2 import io
from . import utilities as _utilities
from .collector import DistributedCollectorNode as _CollectorRuntime
from ..utils import transfer_codec as _transfer_codec
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
from .image_list import DistributedImageListEnd, DistributedImageListStart
from .media_collector import DistributedMediaCollector
class DistributedSeed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
max=1125899906842624, force_input=False)],
outputs=[io.Int.Output(display_name='seed')],
accept_all_inputs=True,
)
@classmethod
def execute(cls, seed, is_worker=False, worker_id=''):
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
class DistributedValue(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedValue', display_name='Distributed Value', category='utils',
inputs=[io.String.Input('default_value', default=''),
io.String.Input('worker_values', default='{}')],
outputs=[io.AnyType.Output(display_name='value')],
accept_all_inputs=True,
)
@classmethod
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
return io.NodeOutput(*_utilities.DistributedValue().distribute(
default_value, worker_values, is_worker, worker_id))
class DistributedModelName(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
inputs=[io.String.Input('text', default='')],
outputs=[io.AnyType.Output(display_name='output')],
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
)
@classmethod
def execute(cls, text):
result = _utilities.DistributedModelName().log_input(
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
return io.NodeOutput(*result['result'], ui=result['ui'])
class ImageBatchDivider(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
inputs=[io.Image.Input('images'),
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
display_mode=io.NumberDisplay.number,
tooltip='Number of parts to divide the batch into')],
# The existing frontend still displays only divide_by sockets.
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
is_output_node=True,
)
@classmethod
def execute(cls, images, divide_by):
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
class AudioBatchDivider(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
inputs=[io.Audio.Input('audio'),
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
display_mode=io.NumberDisplay.number,
tooltip='Number of sequential time segments to create')],
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
is_output_node=True,
)
@classmethod
def execute(cls, audio, divide_by):
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
class DistributedEmptyImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
io.Int.Input('width', default=64, min=1, max=4096, step=1),
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, height, width, channels):
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
class DistributedCollector(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
inputs=[io.Boolean.Input('load_balance', default=False,
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,
tooltip=_transfer_codec.TRANSFER_FORMAT_TOOLTIP),
io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
min=1, max=100, step=1, optional=True,
tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
is_input_list=True, accept_all_inputs=True,
)
@classmethod
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
worker_id='', pass_through=False, delegate_only=False,
transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
return io.NodeOutput(*_CollectorRuntime().run(
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
worker_batch_size=worker_batch_size, worker_id=worker_id,
pass_through=pass_through, delegate_only=delegate_only,
transfer_format=transfer_format, transfer_quality=transfer_quality))
class UltimateSDUpscaleDistributed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='UltimateSDUpscaleDistributed',
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
inputs=[
io.Image.Input('upscaled_image'), io.Model.Input('model'),
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
io.Int.Input('steps', default=20, min=1, max=10000),
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
io.Int.Input('padding', default=32, min=0, max=256, step=8),
io.Int.Input('mask_blur', default=8, min=0, max=256),
io.Boolean.Input('force_uniform_tiles', default=True),
io.Boolean.Input('tiled_decode', default=False),
], outputs=[io.Image.Output()], accept_all_inputs=True,
)
@classmethod
def fingerprint_inputs(cls, **kwargs):
return _UpscaleRuntime.IS_CHANGED(**kwargs)
@classmethod
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
return io.NodeOutput(*_UpscaleRuntime().run(
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
class DistributedWorkerSink(io.ComfyNode):
"""Internal output node that terminates pruned worker prompts.
Worker prompts drop everything downstream of the distributed node, which
leaves no output node for ComfyUI to execute towards. This sink provides
one without writing preview files (PreviewImage wrote every frame to the
worker's temp directory). Hidden from search/menus; inserted by the master.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedWorkerSink', display_name='Distributed Worker Sink (internal)', category='utils',
inputs=[io.AnyType.Input('value')],
is_output_node=True, is_dev_only=True,
)
@classmethod
def execute(cls, value):
return io.NodeOutput()
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedMediaCollector]
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.3.1"
version = "1.5.0"
license = {file = "LICENSE"}
dependencies = []
+1
View File
@@ -0,0 +1 @@
"""Internal extension lifecycle support (not a node provider)."""
+38
View File
@@ -0,0 +1,38 @@
"""Initialize routes, distributed state and the existing worker lifecycle."""
import atexit
import os
import server
from ..utils.config import CONFIG_FILE, ensure_config_exists
from ..utils.logging import debug_log
from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
from ..upscale.job_store import ensure_tile_jobs_initialized
_initialized = False
def initialize():
"""Called by ComfyExtension.on_load; initialize once per loaded package."""
global _initialized
if _initialized:
return
ensure_config_exists()
from .. import api # noqa: F401 - registers the existing @routes.* handlers
from ..api.queue_orchestration import ensure_distributed_state
ensure_distributed_state(server.PromptServer.instance)
ensure_tile_jobs_initialized()
from ..image_list import runtime as image_list_runtime
from ..image_list.lifecycle import install_lifecycle
install_lifecycle(server.PromptServer.instance, image_list_runtime)
if not os.environ.get('COMFYUI_IS_WORKER'):
atexit.register(sync_cleanup)
delayed_auto_launch()
register_async_signals()
_initialized = True
debug_log('Loaded Distributed nodes.')
debug_log(f'Config file: {CONFIG_FILE}')
+251
View File
@@ -0,0 +1,251 @@
# Standalone `comfy-media-v1` endpoint bridge
This directory is independently installable. It registers two internal Comfy V3
nodes, `MediaImportV1` and `MediaExportV1`, and provides the asynchronous HTTP
executor `execute_media_request`. It does **not** start ComfyUI, provision an
endpoint, install models, or provide a provider worker/supervisor. The endpoint
needs this bridge, not the full ComfyUI-Distributed extension.
For collector setup, list ordering, supported native VIDEO states, and PR2
limitations, see [typed serverless media](../docs/serverless-media.md).
## Install in the endpoint image
1. Copy this complete directory to `<ComfyUI>/custom_nodes/serverless_bridge`.
Use a compatible ComfyUI with `comfy_api.v0_0_2` and the native VIDEO classes
used by [video_codec.py](video_codec.py). These exact-class/private-state
adapters are tied to the supported core implementation, not arbitrary VIDEO
plugins or future core layouts.
2. Restart ComfyUI so `comfy_entrypoint()` registers the two nodes. Keep the
workflow's own models and custom nodes installed separately.
3. Make `<ComfyUI>/custom_nodes` importable by the handler process (for example,
via its `PYTHONPATH`). The handler uses `aiohttp`; the nodes/codecs use
ComfyUI's PyTorch/native media dependencies. The handler module itself does
not import ComfyUI, PyTorch, Runpod, or Distributed bootstrap code.
4. Give the handler filesystem access to **the same input directory** used by
that ComfyUI instance, plus its trusted HTTP URL. These are operator settings,
never request-controlled paths/URLs. Staging/cleanup requires POSIX directory
descriptor and no-follow file operations.
Installing the nodes alone is insufficient: the endpoint must route media-v1
inputs through this executor and implement the uncertainty policy below. A
legacy image handler cannot interpret the typed envelope.
## Callable integration
```python
from serverless_bridge.handler import execute_media_request
async def execute_endpoint_input(
input_data, *, comfy_url, input_dir, temp_dir=None,
download=None, publish=None,
):
# Invoke only under a supervisor implementing the uncertainty policy below.
# Exceptions intentionally propagate; this is NOT a provider worker loop.
return await execute_media_request(
input_data,
comfy_url=comfy_url,
input_dir=input_dir,
timeout_seconds=1800,
session=None,
download=download,
publish=publish,
temp_dir=temp_dir,
)
```
Pass the provider job's **`input` object**, not the whole job. Return the resulting
`{"media": ...}` object as provider output; Runpod exposes it as `output.media`.
Do not add another `output` wrapper inside it. `timeout_seconds` must be finite
and positive. If supplied, `session` is a caller-owned `aiohttp.ClientSession`
for local Comfy HTTP and remains open; otherwise the executor creates/closes its
own session. Input dictionaries are copied before graph rewriting.
**This callable is not a turnkey Runpod handler.** An outer SDK that catches every
exception and immediately reuses the worker is unsafe. The supervisor must
recognize `serverless_bridge.handler.MediaExecutionUncertain`, stop intake, and
quarantine/reconcile or terminate the owning environment before reuse. Merely
returning `{"error": ...}` or marking a provider job `FAILED` is not that policy.
## Request and response
A minimal text request builder (no files or model required by this graph):
```python
import json
import uuid
def make_text_request(text="hello"):
request_id = uuid.uuid4().hex
binding = {
"binding_id": "text-binding",
"collector_id": "collector-1",
"job_id": "collector-job-1",
"participant_id": "serverless:example:0",
"node_id": "export-1",
"socket": "texts",
"kind": "STRING",
}
return {
"workflow": {
"text-1": {"class_type": "PrimitiveString", "inputs": {"value": text}},
"export-1": {
"class_type": "MediaExportV1",
"inputs": {
"request_id": request_id,
"binding_json": json.dumps(binding),
"texts": ["text-1", 0],
},
},
},
"images": [],
"media": {
"version": 1,
"profile": "comfy-media-v1",
"request_id": request_id,
"artifacts": [],
"inputs": [],
"outputs": [binding],
},
}
```
`workflow` is an API graph, not an editor workflow. Top-level input fields are
exactly `workflow`, `media`, and optional `images`. `media` has exactly the six
fields shown above. `media.outputs` declares one exporter per connected socket;
its binding must match that node's `binding_json` and request ID exactly. Socket
names/types are `images/IMAGE`, `audio/AUDIO`, `masks/MASK`, `texts/STRING`, and
`video/VIDEO`. Use unique request IDs. Identifiers are 1–200 characters from
`[A-Za-z0-9_.:/-]`; request scopes additionally reject empty, `.` or `..` slash
components. The executor validates the graph and submits **once**, with exactly
the declared exporter IDs as `partial_execution_targets`. Unrelated graph roots,
cycles, undeclared exporters, or caller-supplied `MediaImportV1` nodes fail.
For typed inputs, `media.artifacts` contains records with exactly:
| Field | Value |
| --- | --- |
| `id`, `kind` | Unique artifact ID and one of the five native types |
| `descriptor` | Kind-specific codec metadata; use the shared codec, not pickle |
| `byte_length`, `sha256` | Exact payload byte count and lowercase SHA-256 hex |
| `transport` | `{"type": "inline", "data_base64": "..."}` or the URL transport below |
Each `media.inputs` entry is
`{"node_id": "consumer", "input_name": "socket", "artifact_id": "artifact-id"}`.
It replaces an existing consumer input with a generated import link; every
artifact must be used, and destinations cannot repeat. The executor verifies
bytes/checksums, stages owned files, inserts imports, then prunes disconnected
loader ancestors. Optional legacy `images` entries are `{ "name": "input.png",
"image": "<base64>" }`, must match stock `LoadImage` references, and are rewritten
to owned staging paths. They are not an arbitrary filename-upload interface.
The returned `media` object has exactly `version`, `profile`, `request_id`,
`bindings`, `artifacts`, and `items`:
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
- `artifacts`: the same artifact schema above with inline or URL transport.
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
records covering every index exactly once, ordered by `item_index` when decoded.
All identity, completeness, type, size, and checksum checks must pass; do not
consume a partial result. The internal history transport `comfy-file` is local
only and must not be returned to the master. The executor reads validated
`media_exports` history records and bounded `/view` payloads before constructing
the external envelope. [contract.py](contract.py) is the schema authority.
## Optional image return encoding
IMAGE `MediaExportV1` nodes may also contain literal `transfer_format` and
`transfer_quality` controls. Defaults are `exact` and `90`; old exact requests
omit the controls and retain their original representation. Other media exports
must not carry image controls. The master compiler supplies them from the typed
Collector's widgets; callers should not alter binding identity metadata.
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`, `h265`, `av1`.
`exact` preserves float tensor bits. The image/video codecs quantize RGB/grayscale
to 8-bit; PNG/FFV1 are lossless at that precision, while the other codecs are
lossy. Quality is an integer 1–100 and affects only lossy formats. IMAGE batch
order, shape and channel count are retained, and alpha is not discarded.
Native VIDEO/MASK/AUDIO/STRING still use their original exact codecs.
Copy/update the complete bridge directory on both sides, including
`image_transfer.py` and `image_codec.py`. Pillow and the selected CPU/software
PyAV codec must be available; missing encoders fail without automatic fallback.
Compressed IMAGE artifacts use `image-transfer-v1`; decode through the shared
`codec.decode_payload`, which enforces payload and decoded-shape limits.
See [image transfer choices](../docs/serverless-media.md#image-return-format-and-quality).
## Transport budgets and ownership
- Contract maxima: 256 MiB per artifact, 512 MiB aggregate artifact bytes,
128 artifacts, 64 items per binding. Request input totals include decoded
legacy-image bytes. Decoders validate tensor/VIDEO sizes and layouts before
allocating native storage; these are payload limits, **not peak-RSS guarantees**.
- The built-in master currently emits inline inputs only. Its complete Runpod
request JSON is capped at **10,000,000 bytes**, including base64, graph and
policy. The larger contract limits do not bypass this cap; there is no
built-in large-input/object-storage uploader.
- The executor's result JSON is capped at **10,000,000 bytes**. Inline results
are the default. When needed, `publish` is called for selected larger
artifacts to fit that budget; supplying it does not force all results to URLs.
Without a publisher, an oversized response fails instead of being truncated.
- Optional **async** `download(url, *, max_bytes)` returns `bytes`. The executor
supplies the declared artifact length and checks exact length and SHA-256.
The callback must enforce public HTTPS, DNS/IP restrictions, redirect checks,
and bounded streaming **before allocation**; syntactic contract validation
does not provide SSRF protection. No downloader is installed by default.
- Optional **async** `publish(metadata, payload)` uploads the exact `bytes` and
returns exactly `{"type": "url", "url": "https://...", "expires_at": "..."}`.
`metadata` is the artifact record without `transport`. Expiry must be a future,
timezone-aware ISO timestamp; URLs use HTTPS port 443, no userinfo/fragment,
and at most 8192 characters. Keep objects available long enough for collection.
The callback/operator owns upload credentials, object expiry and cleanup,
including orphan uploads after failure. There is no bridge storage backend.
- The master uses a separate bounded public-HTTPS downloader without provider
credentials. Never send provider credentials to artifact hosts or log signed
URLs. Endpoint callbacks must enforce equivalent protections themselves.
Callbacks must cooperate with async cancellation, finish/abort their own I/O,
and not detach readers/uploads that outlive the callback. Synchronous code that
blocks the event loop cannot be hard-timed-out by this executor.
## Cancellation, quarantine, and file retention
`timeout_seconds` covers staging, submission, polling, downloads and publication.
On failure/cancellation, there is a separate **10-second callback-drain grace**
and, when needed, a **10-second prompt-reconciliation grace**. These are recovery
allowances, not evidence that execution stopped. Ambiguous `/prompt` submission
is never retried. An identified pending prompt may be deleted from `/queue`,
with exact read-back; running prompts are never sent `/interrupt`. The pinned
core route has a snapshot/global-interrupt race even when given a prompt ID.
`MediaExecutionUncertain` means prompt acceptance/termination **or callback
teardown** is unconfirmed. Even successful terminal Comfy history is insufficient
if an async downloader/publisher ignores cancellation. The bridge retains input
and exporter files in that case; late callback exit does not automatically
remove retained files. Before accepting another job on that environment:
1. Quarantine it and stop new intake; do not automatically resubmit the request.
2. Reconcile the local prompt and all readers, or terminate the owning execution
environment when their termination cannot be established.
3. Only then clean retained files and allow reuse/replacement. A provider terminal
job status alone cannot certify a surviving local Comfy process or callback.
Staged inputs use random exclusive files under
`input_dir/distributed-media-v1/<sha256(request_id)>/`. Cleanup removes only files
created by that invocation, after confirmed terminal state (or before submission)
and reader teardown; it never recursively deletes a shared namespace.
`temp_dir=None` leaves exporter retention to the Comfy temp-directory/environment
lifecycle. Optionally pass the **trusted local Comfy temp root** to enable
best-effort cleanup of validated, invocation-owned exporter UUID `.bin` files
under the corresponding hashed namespace, only after terminal history and
collection/publisher teardown. Symlinked, hard-linked, replaced, missing, or
unidentifiable files are not claimed; directories and all unidentified leftovers
remain lifecycle-owned. This is not a temp-directory sweep. Independently decoded
VIDEO objects do not depend on either staging or exporter files.
Verification scope: source and offline tests. No live provider/GPU validation is
claimed, and this directory is not a tested turnkey provider deployment.
+14
View File
@@ -0,0 +1,14 @@
"""Standalone typed-media bridge; copy this directory into endpoint custom_nodes.
Importing codecs does not load ComfyUI, providers or Distributed bootstrap.
"""
async def comfy_entrypoint():
from comfy_api.v0_0_2 import ComfyExtension
from .nodes import MediaImportV1, MediaExportV1
class MediaBridgeExtension(ComfyExtension):
async def get_node_list(self):
return [MediaImportV1, MediaExportV1]
return MediaBridgeExtension()
+189
View File
@@ -0,0 +1,189 @@
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
little-endian floating-point values, never a pickle or an image/audio container.
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
"""
import ctypes
import sys
import torch
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
_DTYPES = {
"float16": (torch.float16, 2),
"bfloat16": (torch.bfloat16, 2),
"float32": (torch.float32, 4),
"float64": (torch.float64, 8),
}
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
_LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
def _validate_kind(kind):
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
raise ValueError("Unsupported payload kind")
def _byte_limit(max_bytes):
if type(max_bytes) is not int or max_bytes < 0:
raise ValueError("max_bytes must be a nonnegative integer")
return min(max_bytes, MAX_PAYLOAD_BYTES)
def _sample_rate(value):
if type(value) is not int or value <= 0:
raise ValueError("AUDIO sample_rate must be a positive integer")
return value
def _tensor_size(kind, shape, itemsize, limit):
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
raise ValueError("Tensor shape has an invalid rank")
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
raise ValueError("Tensor shape dimensions must be positive integers")
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
size = itemsize
for dimension in shape:
# Divide before multiplying, including for arbitrary-size Python ints.
if dimension > limit // size:
raise ValueError("Tensor payload exceeds byte limit")
size *= dimension
return size
def _native_little_endian(payload, itemsize):
"""Convert native to wire byte order, or vice versa (an involution)."""
if sys.byteorder == "little":
return payload
swapped = bytearray(len(payload))
for offset in range(itemsize):
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
return bytes(swapped)
def _encode_string(value, limit):
if not isinstance(value, str):
raise ValueError("STRING value must be text")
# Check the exact UTF-8 size before allocating the encoded payload.
if len(value) > limit:
raise ValueError("STRING payload exceeds byte limit")
size = 0
for character in value:
codepoint = ord(character)
if 0xD800 <= codepoint <= 0xDFFF:
raise ValueError("STRING value must be valid UTF-8 text")
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
if size > limit:
raise ValueError("STRING payload exceeds byte limit")
return {"encoding": "utf8-v1"}, value.encode("utf-8")
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
"""Encode a typed value without dtype conversion, normalization or GPU use.
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
plus a positive integer sample_rate. STRING accepts exact Unicode text.
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
"""
_validate_kind(kind)
limit = _byte_limit(max_bytes)
if kind == "STRING":
return _encode_string(value, limit)
sample_rate = None
if kind == "AUDIO":
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
raise ValueError("AUDIO value must contain waveform and sample_rate")
sample_rate = _sample_rate(value["sample_rate"])
value = value["waveform"]
if not isinstance(value, torch.Tensor):
raise ValueError("Tensor value must be a torch.Tensor")
if value.device.type != "cpu":
raise ValueError("Tensor value must be on CPU")
if value.layout != torch.strided or value.is_nested:
raise ValueError("Tensor value must have a dense strided layout")
if value.dtype not in _DTYPE_NAMES:
raise ValueError("Unsupported tensor dtype")
dtype_name = _DTYPE_NAMES[value.dtype]
_, itemsize = _DTYPES[dtype_name]
shape = list(value.shape)
size = _tensor_size(kind, shape, itemsize, limit)
tensor = value.detach().resolve_neg().contiguous()
if not torch.isfinite(tensor).all().item():
raise ValueError("Tensor values must be finite")
# A uint16 view preserves BF16 bits without conversion through float32.
if tensor.dtype == torch.bfloat16:
tensor = tensor.view(torch.uint16)
# The live, contiguous CPU tensor owns this pointer for the entire copy.
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
descriptor = {
"encoding": "tensor-le-v1", "dtype": dtype_name,
"shape": shape, "layout": _LAYOUTS[kind],
}
if kind == "AUDIO":
descriptor["sample_rate"] = sample_rate
return descriptor, payload
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
All metadata, shape products and byte lengths are checked before tensor
allocation. No GPU device discovery, transfers or object deserialization
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
"""
_validate_kind(kind)
limit = _byte_limit(max_bytes)
if not isinstance(descriptor, dict):
raise ValueError("Payload descriptor must be a dictionary")
if not isinstance(payload, bytes):
raise ValueError("Payload must be bytes")
if len(payload) > limit:
raise ValueError("Payload exceeds byte limit")
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
from .image_codec import decode_image
return decode_image(descriptor, payload, max_bytes=limit)
if kind == "STRING":
if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1":
raise ValueError("Invalid STRING descriptor encoding or fields")
try:
return payload.decode("utf-8")
except UnicodeDecodeError:
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
fields = {"encoding", "dtype", "shape", "layout"}
if kind == "AUDIO":
fields.add("sample_rate")
if set(descriptor) != fields:
raise ValueError("Invalid tensor descriptor fields")
if descriptor["encoding"] != "tensor-le-v1":
raise ValueError("Unsupported tensor descriptor encoding")
if descriptor["layout"] != _LAYOUTS[kind]:
raise ValueError("Invalid tensor descriptor layout")
dtype_name = descriptor["dtype"]
if type(dtype_name) is not str or dtype_name not in _DTYPES:
raise ValueError("Unsupported tensor descriptor dtype")
dtype, itemsize = _DTYPES[dtype_name]
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
if len(payload) != size:
raise ValueError("Tensor payload length does not match shape and dtype")
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
# buffer, then clone so the returned tensor owns its storage independently.
buffer = bytearray(_native_little_endian(payload, itemsize))
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
if not torch.isfinite(tensor).all().item():
raise ValueError("Tensor payload values must be finite")
tensor = tensor.clone().reshape(descriptor["shape"])
if kind == "AUDIO":
return {"waveform": tensor, "sample_rate": sample_rate}
return tensor
+249
View File
@@ -0,0 +1,249 @@
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
"""
import base64
import binascii
import hashlib
import json
import re
from datetime import datetime, timezone
from urllib.parse import urlsplit
PROFILE = "comfy-media-v1"
VERSION = 1
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
MAX_TOTAL_BYTES = 512 * 1024 * 1024
MAX_ARTIFACTS = 128
MAX_ITEMS = 64
MAX_DESCRIPTOR_BYTES = 128 * 1024
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
_HASH = re.compile(r"[a-f0-9]{64}\Z")
def identifier(value):
if not isinstance(value, str) or not _ID.fullmatch(value):
raise ValueError("Invalid media identifier")
return value
def bounded_int(value, maximum, *, minimum=0):
if type(value) is not int or not minimum <= value <= maximum:
raise ValueError("Invalid media count or size")
return value
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
def walk(item, depth=0):
if depth > 12:
raise ValueError("Media metadata nesting exceeds limit")
if isinstance(item, dict):
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
raise ValueError("Invalid media metadata object")
for child in item.values():
walk(child, depth + 1)
elif isinstance(item, list):
if len(item) > 4096:
raise ValueError("Media metadata list exceeds limit")
for child in item:
walk(child, depth + 1)
elif item is not None and type(item) not in (str, int, float, bool):
raise ValueError("Invalid media metadata value")
walk(value)
try:
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
except (ValueError, TypeError, UnicodeError, RecursionError):
raise ValueError("Invalid media JSON") from None
if len(raw) > max_bytes:
raise ValueError("Media JSON exceeds byte limit")
return raw
def validate_binding(binding):
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
raise ValueError("Invalid media binding")
for field in BINDING_FIELDS - {"socket", "kind"}:
identifier(binding[field])
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
raise ValueError("Media socket/type mismatch")
return binding
def validate_bindings(bindings):
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
raise ValueError("Invalid media bindings count")
result = {}
destinations = set()
for binding in bindings:
validate_binding(binding)
bid = binding["binding_id"]
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
if bid in result or destination in destinations:
raise ValueError("Duplicate media binding")
result[bid] = binding
destinations.add(destination)
return result
def _validate_url_transport(transport):
if set(transport) != {"type", "url", "expires_at"}:
raise ValueError("Invalid media URL transport")
url = transport["url"]
if not isinstance(url, str) or len(url) > 8192:
raise ValueError("Invalid media URL")
try:
parsed = urlsplit(url)
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
raise ValueError()
if parsed.port not in (None, 443):
raise ValueError()
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
raise ValueError()
except (ValueError, TypeError, AttributeError):
raise ValueError("Invalid or expired media URL") from None
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid media artifact")
identifier(artifact["id"])
if artifact["kind"] not in SOCKET_KINDS.values():
raise ValueError("Unsupported media kind")
if not isinstance(artifact["descriptor"], dict):
raise ValueError("Invalid media descriptor")
validate_json(artifact["descriptor"])
size = bounded_int(artifact["byte_length"], max_bytes)
digest = artifact["sha256"]
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
raise ValueError("Invalid media checksum")
transport = artifact["transport"]
if not isinstance(transport, dict):
raise ValueError("Invalid media transport")
if transport.get("type") == "inline":
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
raise ValueError("Invalid inline media transport")
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
raise ValueError("Inline media length mismatch")
elif transport.get("type") == "url":
_validate_url_transport(transport)
else:
raise ValueError("Unsupported media transport")
return artifact
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
if not isinstance(payload, bytes):
raise ValueError("Media payload must be bytes")
bounded_int(len(payload), max_bytes)
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
"transport": transport if transport is not None else {
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
return validate_artifact(result, max_bytes=max_bytes)
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
validate_artifact(artifact, max_bytes=max_bytes)
transport = artifact["transport"]
if transport["type"] == "inline":
try:
payload = base64.b64decode(transport["data_base64"], validate=True)
except (ValueError, binascii.Error):
raise ValueError("Invalid media base64") from None
else:
if download is None:
raise ValueError("Media URL download is not configured")
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
raise ValueError("Media byte length mismatch")
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Media checksum mismatch")
return payload
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
"""Validate ALL identity/completeness before downloading or delivering anything."""
expected = validate_bindings(bindings)
identifier(request_id)
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
raise ValueError("Invalid media response envelope")
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
raise ValueError("Unsupported media response version/profile")
if response["request_id"] != request_id:
raise ValueError("Media request identity mismatch")
returned = response["bindings"]
if not isinstance(returned, list) or len(returned) != len(expected):
raise ValueError("Missing or extra media bindings")
counts = {}
for row in returned:
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
raise ValueError("Invalid result binding")
binding = {key: row[key] for key in BINDING_FIELDS}
validate_binding(binding)
bid = binding["binding_id"]
if bid in counts or expected.get(bid) != binding:
raise ValueError("Media binding identity mismatch")
# Connected outputs must produce at least one value. Empty STRING is a value.
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
artifacts = response["artifacts"]
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
raise ValueError("Invalid media artifacts count")
indexed = {}
total = 0
for artifact in artifacts:
validate_artifact(artifact)
aid = artifact["id"]
if aid in indexed:
raise ValueError("Duplicate media artifact")
total += artifact["byte_length"]
if total > max_total_bytes:
raise ValueError("Media response exceeds total byte limit")
indexed[aid] = artifact
items = response["items"]
if not isinstance(items, list) or len(items) != sum(counts.values()):
raise ValueError("Missing or extra media items")
ordered = {bid: [None] * count for bid, count in counts.items()}
used = set()
for item in items:
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
raise ValueError("Invalid media item")
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
if bid not in counts or aid not in indexed:
raise ValueError("Unknown media item identity")
index = bounded_int(item["item_index"], counts[bid] - 1)
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
raise ValueError("Duplicate or mistyped media item")
ordered[bid][index] = indexed[aid]
used.add(aid)
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
raise ValueError("Incomplete media result manifest")
return ordered
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
# Resolve and decode before exposing a single result to the collector queue.
from .codec import decode_payload
cache = {}
decoded = {}
for bid, artifacts in ordered.items():
values = []
for artifact in artifacts:
aid = artifact["id"]
if aid not in cache:
payload = await resolve_artifact(artifact, download=download)
if artifact["kind"] == "VIDEO":
from .video_codec import decode_video
cache[aid] = decode_video(artifact["descriptor"], payload)
else:
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
values.append(cache[aid])
decoded[bid] = values
return decoded
+794
View File
@@ -0,0 +1,794 @@
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
Only identified invocation-owned files are deleted. Optional trusted local
``temp_dir`` enables exporter cleanup after terminal history and collection.
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
The pinned native /interrupt route is not atomic: running-prompt cancellation
fails closed and retains files instead of risking an unrelated successor prompt.
"""
import asyncio
import base64
import binascii
import copy
import hashlib
import json
import math
import os
import re
import stat
import uuid
from pathlib import Path
from urllib.parse import quote, urlsplit
import aiohttp
from .contract import (
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
validate_artifact, validate_bindings, validate_json, validate_response,
)
MAX_RESPONSE_JSON_BYTES = 10_000_000
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
POLL_INTERVAL_SECONDS = 0.25
CANCEL_GRACE_SECONDS = 10.0
DRAIN_GRACE_SECONDS = 10.0
HTTP_TIMEOUT_SECONDS = 15.0
_UNCERTAIN_WORK = set()
GET_ATTEMPTS = 3
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
# Used ONLY to run the shared manifest validator before /view. Never fetched.
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
"expires_at": "9999-01-01T00:00:00Z"}
class MediaExecutionUncertain(RuntimeError):
"""The local prompt may still execute; staged files were deliberately retained."""
class _HTTPError(RuntimeError):
def __init__(self, status):
super().__init__(f"Local Comfy HTTP request failed (status {status})")
self.status = status
def _scope_id(value):
identifier(value)
if any(part in ("", ".", "..") for part in value.split("/")):
raise ValueError("Unsafe media scope identifier")
return value
def _json_bytes(value):
return json.dumps(value, allow_nan=False).encode("utf-8")
def _link(value, workflow):
if (not isinstance(value, list) or len(value) != 2 or
not isinstance(value[0], str) or value[0] not in workflow or
type(value[1]) is not int or value[1] < 0):
raise ValueError("Invalid workflow socket link")
return value[0]
def _validate_request(input_data):
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
or set(input_data) - {"workflow", "media", "images"}):
raise ValueError("Invalid media request envelope")
media = input_data["media"]
if (not isinstance(media, dict) or set(media) !=
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
raise ValueError("Invalid media request controls")
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
raise ValueError("Unsupported media request version/profile")
_scope_id(media["request_id"])
bindings = validate_bindings(media["outputs"])
targets = []
for binding in bindings.values():
for field in BINDING_FIELDS - {"socket", "kind"}:
_scope_id(binding[field])
targets.append(binding["node_id"])
if len(set(targets)) != len(targets):
raise ValueError("Duplicate media exporter node")
workflow = input_data["workflow"]
if not isinstance(workflow, dict) or not workflow:
raise ValueError("Invalid API workflow")
validate_json(workflow, max_bytes=16 * 1024 * 1024)
dependencies = {}
for node_id, node in workflow.items():
_scope_id(node_id)
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
raise ValueError("Invalid API workflow node")
if node["class_type"] == "MediaImportV1":
raise ValueError("Request may not supply its own media import filenames")
if node["class_type"] == "MediaExportV1" and node_id not in targets:
raise ValueError("Undeclared media exporter")
dependencies[node_id] = set()
for name, value in node["inputs"].items():
_scope_id(name)
if isinstance(value, list):
dependencies[node_id].add(_link(value, workflow))
for binding in bindings.values():
node = workflow.get(binding["node_id"])
if node is None or node["class_type"] != "MediaExportV1":
raise ValueError("Missing declared MediaExportV1 node")
controls = node["inputs"]
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
required = {"request_id", "binding_json", binding["socket"]}
optional = {"transfer_format", "transfer_quality"} if binding["kind"] == "IMAGE" else set()
if not required <= set(controls) or set(controls) - required - optional:
raise ValueError("Media exporter must have exactly its declared socket and controls")
validate_options(controls.get("transfer_format", DEFAULT_FORMAT),
controls.get("transfer_quality", DEFAULT_QUALITY))
if controls["request_id"] != media["request_id"]:
raise ValueError("Media exporter request identity mismatch")
text = controls["binding_json"]
if not isinstance(text, str) or len(text) > 128 * 1024:
raise ValueError("Invalid media exporter binding JSON")
try:
declared = json.loads(text)
except (ValueError, RecursionError):
raise ValueError("Invalid media exporter binding JSON") from None
if declared != binding:
raise ValueError("Media exporter binding identity mismatch")
_link(controls[binding["socket"]], workflow)
# The execution roots are exactly the declared exporters, not every class
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
visited = set()
pending = list(targets)
while pending:
node_id = pending.pop()
if node_id not in visited:
visited.add(node_id)
pending.extend(dependencies[node_id])
if visited != set(workflow):
raise ValueError("Workflow contains nodes outside declared output roots")
if any(set(targets) & sources for sources in dependencies.values()):
raise ValueError("A media exporter cannot be an upstream dependency")
# Kahn traversal bounds cycle checking without relying on recursion depth.
counts = {node: len(sources) for node, sources in dependencies.items()}
consumers = {node: [] for node in workflow}
for node, sources in dependencies.items():
for source in sources:
consumers[source].append(node)
pending = [node for node, count in counts.items() if not count]
processed = 0
while pending:
processed += 1
for consumer in consumers[pending.pop()]:
counts[consumer] -= 1
if not counts[consumer]:
pending.append(consumer)
if processed != len(workflow):
raise ValueError("Workflow contains a dependency cycle")
artifacts = media["artifacts"]
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
raise ValueError("Invalid input artifacts count")
indexed, total = {}, 0
for artifact in artifacts:
validate_artifact(artifact)
aid = _scope_id(artifact["id"])
if aid in indexed:
raise ValueError("Duplicate input artifact")
indexed[aid] = artifact
total += artifact["byte_length"]
inputs = media["inputs"]
if not isinstance(inputs, list) or len(inputs) > 4096:
raise ValueError("Invalid input references")
destinations, used = set(), set()
exporters = {binding["node_id"]: binding for binding in bindings.values()}
for ref in inputs:
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
raise ValueError("Invalid media input reference")
for value in ref.values():
_scope_id(value)
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
or aid not in indexed or (node_id, name) in destinations):
raise ValueError("Unknown or duplicate media input reference")
binding = exporters.get(node_id)
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
raise ValueError("Exporter input must match its declared media socket and kind")
destinations.add((node_id, name))
used.add(aid)
if used != set(indexed):
raise ValueError("Every input artifact must be used")
images = input_data.get("images", [])
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
raise ValueError("Invalid legacy images count")
names = set()
legacy = []
for image in images:
if not isinstance(image, dict) or set(image) != {"name", "image"}:
raise ValueError("Invalid legacy image")
name, encoded = image["name"], image["image"]
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
or name in names or not isinstance(encoded, str)):
raise ValueError("Invalid or duplicate legacy image basename")
names.add(name)
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
raise ValueError("Legacy image exceeds byte limit")
# Exact decoded length, including padding, before allocating bytes.
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
raise ValueError("Invalid legacy image base64 length")
total += size
matches = [node_id for node_id, node in workflow.items()
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
if not matches or any((node_id, "image") in destinations for node_id in matches):
raise ValueError("Unused or conflicting legacy image binding")
legacy.append((encoded, size, matches))
if total > MAX_TOTAL_BYTES:
raise ValueError("Media inputs exceed total byte limit")
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
class _Staging:
"""Exclusive random files under a hashed namespace; never delete siblings."""
def __init__(self, input_dir, request_id):
self.root = Path(input_dir).resolve()
self.key = hashlib.sha256(request_id.encode()).hexdigest()
self.relative = f"distributed-media-v1/{self.key}"
self.names = []
self.fd = None
def _open(self):
self.root.mkdir(parents=True, exist_ok=True)
directory = self.root
for component in ("distributed-media-v1", self.key):
directory = directory / component
if directory.is_symlink():
raise ValueError("Unsafe media staging directory symlink")
directory.mkdir(exist_ok=True, mode=0o700)
if not directory.resolve().is_relative_to(self.root):
raise ValueError("Unsafe media staging directory")
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
def write(self, payload):
if self.fd is None:
self._open()
name = uuid.uuid4().hex + ".bin"
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
self.names.append(name)
with os.fdopen(fd, "wb") as stream:
stream.write(payload)
return f"{self.relative}/{name}"
def close(self, *, cleanup):
if self.fd is not None:
try:
if cleanup:
for name in self.names:
try:
os.unlink(name, dir_fd=self.fd)
except FileNotFoundError:
pass
finally:
os.close(self.fd)
self.fd = None
if cleanup:
# Never recursively remove a shared namespace or another run.
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
try:
directory.rmdir()
except OSError:
pass
class _OwnedExports:
"""Best-effort retirement of validated terminal-history files, never a sweep.
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
files stay with the endpoint lifecycle. Never use history-supplied path strings.
"""
def __init__(self, temp_dir, request_id):
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
self.key = hashlib.sha256(request_id.encode()).hexdigest()
self.fd = None
self.files = {}
def _open(self):
if self.fd is not None:
return True
if self.root is None or ".." in self.root.parts:
return False
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
fd = None
try:
fd = os.open(self.root.anchor, flags)
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
child = os.open(component, flags, dir_fd=fd)
os.close(fd)
fd = child
self.fd, fd = fd, None
return True
except OSError:
return False
finally:
if fd is not None:
os.close(fd)
def _identity(self, name):
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
return None
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
def identify(self, entry, media):
if self.root is None or not isinstance(entry.get("outputs"), dict):
return
# A failed prompt may contain valid completed exporters alongside missing
# or malformed ones. Claim only independently validated declared records.
for binding in media["outputs"]:
node_id = binding["node_id"]
try:
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
dict(media, outputs=[binding]))
except ValueError:
continue
if not files or not self._open():
continue
for artifact in files:
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
try:
identity = self._identity(name)
except OSError:
continue
if identity is not None:
self.files.setdefault(name, identity)
def close(self, *, cleanup):
if self.fd is None:
return
try:
if cleanup:
for name, identity in self.files.items():
try:
if self._identity(name) == identity:
os.unlink(name, dir_fd=self.fd)
except OSError:
pass # Endpoint lifecycle owns inaccessible leftovers.
finally:
os.close(self.fd)
self.fd = None
# Deliberately retain directories: they may belong to concurrent runs.
async def _body(response, maximum):
length = response.content_length
if length is not None and length > maximum:
raise ValueError("Local Comfy response exceeds declared byte length")
body = bytearray()
async for chunk in response.content.iter_chunked(64 * 1024):
if len(body) + len(chunk) > maximum:
raise ValueError("Local Comfy response exceeds declared byte length")
body.extend(chunk)
return bytes(body)
class _Execution:
def __init__(self, session, url, stage, media, targets, deadline, exports):
self.session, self.url, self.stage = session, url, stage
self.media, self.targets = media, targets
self.deadline = deadline
self.run_id = uuid.uuid4().hex
self.prompt_id = None
self.attempted = False
self.terminal = False
self.exports = exports
self.interruption_blocked = False
def remaining(self):
remaining = self.deadline - asyncio.get_running_loop().time()
if remaining <= 0:
raise asyncio.TimeoutError("Media execution deadline expired")
return remaining
async def json_request(self, method, path, *, body=None, permit_rejection=False):
async with self.session.request(method, self.url + path, json=body,
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
allow_redirects=False) as response:
raw = await _body(response, MAX_HTTP_JSON_BYTES)
# Native queue-delete and interrupt routes return an empty 200.
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
return {}
try:
value = json.loads(raw)
except (ValueError, UnicodeError, RecursionError):
if response.status != 200:
raise _HTTPError(response.status) from None
raise ValueError("Invalid local Comfy JSON response") from None
if response.status != 200:
# Native /prompt rejects the entire graph with a 400 error.
# Other failures may hide acceptance and are never retried.
if (permit_rejection and response.status == 400 and isinstance(value, dict)
and value.get("error") and "prompt_id" not in value):
self.terminal = True
raise _HTTPError(response.status)
if not isinstance(value, dict):
raise ValueError("Invalid local Comfy JSON object")
return value
async def get_json(self, path):
for attempt in range(GET_ATTEMPTS):
try:
return await self.json_request("GET", path)
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
raise
if attempt == GET_ATTEMPTS - 1:
raise
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
async def history(self):
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
entry = response.get(self.prompt_id)
if entry is None:
return None
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
raise ValueError("Malformed Comfy prompt history")
status = entry["status"]
messages = status.get("messages", [])
if not isinstance(messages, list):
raise ValueError("Malformed Comfy history messages")
failed = status.get("status_str") in {"error", "interrupted"} or any(
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
for message in messages)
if failed or status.get("completed") is True:
self.terminal = True
self.exports.identify(entry, self.media)
if failed or status.get("status_str") != "success":
raise RuntimeError("Local Comfy prompt execution failed")
return entry
return None
async def submit_and_wait(self, workflow):
body = {"prompt": workflow, "client_id": self.run_id,
"extra_data": {"distributed_media_run_id": self.run_id},
"partial_execution_targets": list(self.targets)}
self.remaining() # CPU staging must not spend past the absolute deadline.
self.attempted = True
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
self.prompt_id = _scope_id(response.get("prompt_id"))
if response.get("node_errors"):
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
while True:
entry = await self.history()
if entry is not None:
return entry
await asyncio.sleep(POLL_INTERVAL_SECONDS)
async def settle(self):
"""Bounded recovery; caller shields this from task cancellation."""
deleted = False
while True:
if self.prompt_id:
try:
await self.history()
except RuntimeError:
if not self.terminal:
raise
if self.terminal:
return
queue = await self.get_json("/queue")
groups = []
for key in ("queue_running", "queue_pending"):
rows = queue.get(key)
if not isinstance(rows, list):
raise ValueError("Malformed Comfy queue")
group = {}
for row in rows:
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
raise ValueError("Malformed Comfy queue entry")
pid = _scope_id(row[1])
if pid in group:
raise ValueError("Duplicate Comfy queue identity")
group[pid] = row[3]
groups.append(group)
running, pending = groups
if self.prompt_id is None:
own = [pid for group in groups for pid, extra in group.items()
if extra.get("distributed_media_run_id") == self.run_id
and extra.get("client_id") == self.run_id]
if len(own) > 1:
raise ValueError("Ambiguous Comfy queue ownership")
if own:
self.prompt_id = own[0]
if self.prompt_id in pending and not deleted:
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
deleted = True
elif self.prompt_id in running:
# Pinned ComfyUI server.py checks a snapshot then globally
# interrupts, outside the queue mutex. Supplying prompt_id does
# NOT close that race. No safe capability negotiation exists on
# this route; do not invoke it, even after our own queue read.
self.interruption_blocked = True
raise RuntimeError("Native prompt-scoped interruption unavailable")
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
self.terminal = True # queue deletion confirmed by exact read-back
return
await asyncio.sleep(POLL_INTERVAL_SECONDS)
async def fetch(self, artifact):
transport = artifact["transport"]
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
for attempt in range(GET_ATTEMPTS):
try:
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
if response.status != 200:
raise _HTTPError(response.status)
payload = await _body(response, artifact["byte_length"])
break
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
raise
if attempt == GET_ATTEMPTS - 1:
raise
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
if len(payload) != artifact["byte_length"]:
raise ValueError("Media output byte length mismatch")
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Media output checksum mismatch")
return payload
def _manifest(entry, media):
outputs = entry.get("outputs")
if not isinstance(outputs, dict):
raise ValueError("Missing media history outputs")
expected_ids = {binding["node_id"] for binding in media["outputs"]}
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
for node_id, value in outputs.items()):
raise ValueError("Unexpected media export history")
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
"bindings": [], "artifacts": [], "items": []}
files = []
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
for binding in media["outputs"]:
node = outputs.get(binding["node_id"], {})
records = node.get("media_exports") if isinstance(node, dict) else None
if not isinstance(records, list) or len(records) != 1:
raise ValueError("Missing or duplicate media export result")
record = records[0]
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
or record["request_id"] != media["request_id"]):
raise ValueError("Media export request identity mismatch")
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
raise ValueError("Invalid exported artifacts count")
converted = []
for artifact in record["artifacts"]:
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
raise ValueError("Invalid exported artifact")
transport = artifact["transport"]
aid = artifact.get("id")
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
or transport != {"type": "comfy-file", "filename": aid + ".bin",
"subfolder": subfolder, "folder_type": "temp"}):
raise ValueError("Unsafe or unowned exported media file")
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
validate_response(partial, request_id=media["request_id"], bindings=[binding])
response["bindings"].append(copy.deepcopy(record["binding"]))
response["artifacts"].extend(converted)
response["items"].extend(copy.deepcopy(record["items"]))
files.extend(copy.deepcopy(record["artifacts"]))
# Validate the whole manifest BEFORE the first /view or publisher callback.
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
return response, files
def _publication_plan(response, files, publish):
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
base_size = len(_json_bytes({"media": skeleton}))
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
predicted = base_size + sum(sizes.values())
published = set()
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
# Reserve enough space for a maximum-size escaped URL, including metadata.
# Small remaining inline values preserve their exact bytes and ordering.
for aid in sorted(sizes, key=sizes.get, reverse=True):
if predicted <= MAX_RESPONSE_JSON_BYTES:
break
if sizes[aid] > 65_536:
published.add(aid)
predicted -= sizes[aid] - 65_536
if predicted > MAX_RESPONSE_JSON_BYTES:
raise ValueError("Media response metadata exceeds JSON limit")
return published
async def _run(execution, workflow, legacy, download, publish):
media = execution.media
imports = {}
for artifact in media["artifacts"]:
execution.remaining()
payload = await resolve_artifact(artifact, download=download)
execution.remaining()
filename = execution.stage.write(payload)
node_id = "distributed-media-import:" + uuid.uuid4().hex
while node_id in workflow:
node_id = "distributed-media-import:" + uuid.uuid4().hex
metadata = {key: value for key, value in artifact.items() if key != "transport"}
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
for ref in media["inputs"]:
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
for encoded, size, matches in legacy:
execution.remaining()
try:
payload = base64.b64decode(encoded, validate=True)
except (ValueError, binascii.Error):
raise ValueError("Invalid legacy image base64") from None
if len(payload) != size:
raise ValueError("Legacy image byte length mismatch")
execution.remaining()
filename = execution.stage.write(payload)
for node_id in matches:
workflow[node_id]["inputs"]["image"] = filename
# Replacements can disconnect file loaders whose original endpoint-local
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
# shared ancestors still needed by another branch must survive.
reachable, pending = set(), list(execution.targets)
while pending:
node_id = pending.pop()
if node_id in reachable:
continue
reachable.add(node_id)
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
if isinstance(value, list))
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
entry = await execution.submit_and_wait(workflow)
response, files = _manifest(entry, media)
published = _publication_plan(response, files, publish)
artifacts = []
for artifact in files:
payload = await execution.fetch(artifact)
transport = None
if artifact["id"] in published:
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
transport = await publish(metadata, payload)
if not isinstance(transport, dict) or transport.get("type") != "url":
raise ValueError("Publisher must return an expiring HTTPS URL transport")
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
payload, transport=transport))
response["artifacts"] = artifacts
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
result = {"media": response}
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
raise ValueError("Media response exceeds JSON byte limit")
execution.remaining()
return result
def _release_uncertain_work(task):
_UNCERTAIN_WORK.discard(task)
try:
task.result()
except BaseException:
pass
async def _drain_work(work):
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
if not work.done():
work.cancel()
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
while not work.done():
remaining = deadline - asyncio.get_running_loop().time()
if remaining <= 0:
# Keep a strong reference until it exits, but never delete its files
# asynchronously: uncertain ownership belongs to endpoint recovery.
_UNCERTAIN_WORK.add(work)
work.add_done_callback(_release_uncertain_work)
return False
try:
# Unlike wait_for, wait's timeout does not wait for cancellation.
await asyncio.wait({work}, timeout=remaining)
except asyncio.CancelledError:
continue
_release_uncertain_work(work)
return True
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
bounded cancellation grace is allowed ONLY to establish the local prompt is
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
and retain input files rather than racing an executing prompt. Never retry an
ambiguous submission, even when the exception is a timeout.
Async callback teardown has its own bounded grace; non-cooperative callbacks
also raise MediaExecutionUncertain and retain files, even after terminal history.
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
never a request-supplied path. Only validated exporter UUID.bin files beneath
distributed-media-v1/sha256(request_id) are retired, after all collection and
publisher work has unwound. Unidentified files remain lifecycle-owned. With
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
The pinned native /interrupt route has a check/global-interrupt race. Running
prompts therefore fail closed with MediaExecutionUncertain, without posting
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
"""
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
or timeout_seconds <= 0):
raise ValueError("Invalid media execution timeout")
if not isinstance(comfy_url, str):
raise ValueError("Invalid local Comfy URL")
try:
parsed = urlsplit(comfy_url)
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
or parsed.password or parsed.query or parsed.fragment):
raise ValueError()
parsed.port
except ValueError:
raise ValueError("Invalid local Comfy URL") from None
deadline = asyncio.get_running_loop().time() + timeout_seconds
workflow, media, legacy, targets = _validate_request(input_data)
stage = _Staging(input_dir, media["request_id"])
exports = _OwnedExports(temp_dir, media["request_id"])
owned_session = session is None
if owned_session:
session = aiohttp.ClientSession()
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
work = None
try:
remaining = execution.remaining()
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
return await asyncio.wait_for(asyncio.shield(work), remaining)
except BaseException as exc:
# Retire files only after readers unwind. Cooperative async callbacks
# normally drain immediately; an uncooperative one requires quarantine.
if work is not None and not await _drain_work(work):
raise MediaExecutionUncertain(
"Media callback teardown unconfirmed; owned files retained; "
"quarantine this endpoint before reuse; do not resubmit automatically"
) from exc
if execution.attempted and not execution.terminal:
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
# Shield the bounded recovery even from repeated task.cancel() calls.
while not recovery.done():
try:
await asyncio.shield(recovery)
except asyncio.CancelledError:
continue
except Exception:
break
try:
recovery.result()
except BaseException:
pass
if not execution.terminal:
raise MediaExecutionUncertain(
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
"do not resubmit this request automatically"
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
if execution.interruption_blocked else "")
) from exc
raise
finally:
try:
try:
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
finally:
stage.close(cleanup=(work is None or work.done()) and
(not execution.attempted or execution.terminal))
finally:
if owned_session:
await session.close()
+634
View File
@@ -0,0 +1,634 @@
"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed.
Exact mode is the existing tensor wire format, unchanged. All other modes clip
and truncate RGB/luma to uint8 (as the Collector does), normalize in float32,
then cast to the source dtype. RGBA always carries an exact tensor alpha
sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here,
not the old Collector's JSON/base64 transport.
The compressed wire format concatenates length-delimited stills or one Matroska
video, followed by optional alpha bytes. No archive, path, URL or pickle is
accepted. Limits cover wire bytes, descriptor structure, output tensor bytes,
and aggregate padded pixel buffers, not process RSS/native codec overhead.
Video preflight reads EBML track dimensions and block counts before creating a
decoder. We extract packets ourselves, avoiding native container stream probing
(which can decode before caller checks). FFmpeg max_pixels also bounds
bitstream-driven resolution changes, including native buffer alignment.
Software codecs only, one thread, no fallback to another encoder/format.
"""
import io
import struct
import zlib
from fractions import Fraction
import torch
from . import codec
from .codec import MAX_PAYLOAD_BYTES
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
MAX_FRAMES = 240
MAX_DIMENSION = 8192
MAX_DESCRIPTOR_BYTES = 64 * 1024
MAX_CONTAINER_ELEMENTS = 16384
_VIDEO_CODECS = {
"ffv1": ("ffv1", None, "bgr0"),
"h264": ("libx264", 51, "yuv420p"),
"h265": ("libx265", 51, "yuv420p"),
"av1": ("libsvtav1", 63, "yuv420p"),
}
_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC",
"h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"}
class _BoundedIO(io.BytesIO):
"""Enforce the final file size on every write, including muxer flush/seek."""
def __init__(self, limit):
super().__init__()
self.limit = limit
def write(self, data):
if self.tell() + len(data) > self.limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return super().write(data)
def seek(self, offset, whence=0):
if whence == 0:
position = offset
elif whence == 1:
position = self.tell() + offset
elif whence == 2:
position = self.getbuffer().nbytes + offset
else:
raise ValueError("Invalid image buffer seek")
if not 0 <= position <= self.limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return super().seek(offset, whence)
def _quality_to_crf(quality, maximum):
return int(round(maximum * (1.0 - quality / 100.0) ** 0.6))
def _padded_size(fmt, height, width):
if fmt in ("h264", "h265", "av1"):
height, width = height + height % 2, width + width % 2
# libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures.
if fmt in ("h265", "av1"):
height, width = max(64, height), max(64, width)
return height, width
def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None):
size = codec._tensor_size("IMAGE", shape, itemsize, limit)
count, height, width, _ = shape
if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION:
raise ValueError("IMAGE frame count or dimensions exceed limit")
ph, pw = _padded_size(fmt, height, width)
# Reserve output and padded RGBA-sized decoded buffers BEFORE allocation.
bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw)
# Slice/OBU counts conservatively bound native frames, including AV1 hidden
# frames and multiple outputs from one packet. Reserve 16 bytes per pixel
# for hostile high-bit-depth streams, not only our usual 8-bit RGB frames.
native_count = max(count, coded_frames or count)
if native_count > MAX_FRAMES * 4:
raise ValueError("IMAGE coded frame count exceeds limit")
pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit)
if size + pixels + alpha_size > limit:
raise ValueError("IMAGE decoded aggregate exceeds byte limit")
return ph, pw
def _descriptor_budget(value):
"""Bound the JSON-shaped tree without serializing attacker-sized values."""
remaining = MAX_DESCRIPTOR_BYTES
def walk(obj, depth):
nonlocal remaining
if depth > 5:
raise ValueError("IMAGE descriptor nesting exceeds limit")
remaining -= 8
if remaining < 0:
raise ValueError("IMAGE descriptor exceeds byte limit")
if obj is None:
return
if type(obj) is str:
remaining -= len(obj) * 6
elif type(obj) is int:
if obj.bit_length() > 63:
raise ValueError("IMAGE descriptor integer exceeds limit")
remaining -= 20
elif type(obj) is list:
if len(obj) > MAX_FRAMES:
raise ValueError("IMAGE descriptor list exceeds limit")
for item in obj:
walk(item, depth + 1)
elif type(obj) is dict:
if len(obj) > 10:
raise ValueError("IMAGE descriptor fields exceed limit")
for key, item in obj.items():
if type(key) is not str:
raise ValueError("Invalid IMAGE descriptor key")
walk(key, depth + 1)
walk(item, depth + 1)
else:
raise ValueError("Invalid IMAGE descriptor value")
if remaining < 0:
raise ValueError("IMAGE descriptor exceeds byte limit")
walk(value, 0)
def _alpha_metadata(alpha, shape, limit):
if shape[-1] != 4:
if alpha is not None:
raise ValueError("Unexpected IMAGE alpha sidecar")
return 0, torch.float32
if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}:
raise ValueError("Missing or invalid IMAGE alpha sidecar")
d = alpha["descriptor"]
if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}:
raise ValueError("Invalid IMAGE alpha descriptor")
if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]:
raise ValueError("Invalid IMAGE alpha shape or encoding")
if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES:
raise ValueError("Invalid IMAGE alpha dtype")
dtype, itemsize = codec._DTYPES[d["dtype"]]
size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit)
if type(alpha["size"]) is not int or alpha["size"] != size:
raise ValueError("Invalid IMAGE alpha size")
return size, torch.float64 if dtype == torch.float64 else torch.float32
def _import_av():
try:
import av
except ImportError:
raise ValueError("PyAV is required for IMAGE video transfer") from None
return av
def _encode_still(frame, fmt, quality, limit):
from PIL import Image
sink = _BoundedIO(limit)
if frame.shape[-1] == 1:
frame = frame[..., 0]
image = Image.fromarray(frame)
if fmt in ("legacy_png", "png"):
image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4)
elif fmt == "jpeg":
image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
else:
image.save(sink, format="WEBP", quality=quality, method=4)
return sink.getvalue()
def _encode_video(frames, fmt, quality, height, width, limit):
import numpy as np
av = _import_av()
encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt]
try:
av.codec.Codec(encoder, "w")
except Exception:
raise ValueError("IMAGE software encoder unavailable: " + encoder) from None
sink = _BoundedIO(limit)
options = {}
if maximum is not None:
options["crf"] = str(_quality_to_crf(quality, maximum))
if fmt == "h264":
options["preset"] = "medium"
elif fmt == "h265":
options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"})
elif fmt == "av1":
options.update(preset="8", **{"svtav1-params": "lp=1"})
else:
options.update(level="3", slicecrc="1")
try:
with av.open(sink, mode="w", format="matroska") as container:
stream = container.add_stream(encoder, rate=24)
stream.width, stream.height = width, height
stream.pix_fmt = pix_fmt
stream.codec_context.thread_count = 1
stream.options = options
for index, frame in enumerate(frames):
if frame.shape[-1] == 1:
frame = np.repeat(frame, 3, axis=-1)
h, w = frame.shape[:2]
frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
vf = av.VideoFrame.from_ndarray(frame, format="rgb24")
vf.pts = index
vf.time_base = Fraction(1, 24)
for packet in stream.encode(vf):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
data = sink.getvalue()
if len(data) > limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return data
except av.error.FFmpegError:
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
except ValueError:
raise
except Exception:
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY,
max_bytes=MAX_PAYLOAD_BYTES):
"""Return (descriptor, bytes); reject unsupported inputs, never fall back."""
fmt, quality = validate_options(transfer_format, transfer_quality)
if fmt == "exact":
return codec.encode_payload("IMAGE", value, max_bytes=max_bytes)
limit = codec._byte_limit(max_bytes)
if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
raise ValueError("IMAGE must be a CPU tensor")
if value.layout != torch.strided or value.is_nested:
raise ValueError("IMAGE must have dense strided layout")
if value.dtype not in codec._DTYPE_NAMES:
raise ValueError("Unsupported IMAGE dtype")
shape = list(value.shape)
codec._tensor_size("IMAGE", shape, value.element_size(), limit)
alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0
itemsize = value.element_size()
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
if not torch.isfinite(value).all().item():
raise ValueError("IMAGE values must be finite")
value = value.detach().resolve_neg()
alpha, alpha_payload = None, b""
if shape[-1] == 4:
ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit)
alpha = {"descriptor": ad, "size": len(alpha_payload)}
import numpy as np
def frames():
for frame in value:
arr = frame[..., :3].float().numpy()
with np.errstate(over="ignore", invalid="ignore"):
yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
budget = limit - len(alpha_payload)
if fmt in _VIDEO_CODECS:
parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)]
header = _video_header(parts[0], fmt, ph, pw, shape[0])
coded_frames = _coded_frame_bound(parts[0], fmt, header)
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
else:
parts = []
for frame in frames():
part = _encode_still(frame, fmt, quality, budget)
budget -= len(part)
parts.append(part)
descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality,
"shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype],
"parts": [len(p) for p in parts], "alpha": alpha}
_descriptor_budget(descriptor)
payload = b"".join(parts) + alpha_payload
if len(payload) > limit:
raise ValueError("Encoded IMAGE exceeds byte limit")
return descriptor, payload
def _still_header(data, fmt):
"""Read dimensions before Pillow (WebP may allocate pixels during open)."""
if fmt in ("png", "legacy_png"):
if data[:8] != b"\x89PNG\r\n\x1a\n":
raise ValueError("Invalid IMAGE PNG header")
offset, dimensions, saw_data, chunks = 8, None, False, 0
while offset < len(data):
chunks += 1
if chunks > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE PNG chunk count exceeds limit")
if offset + 12 > len(data):
raise ValueError("Truncated IMAGE PNG chunk")
length = int.from_bytes(data[offset:offset + 4], "big")
end = offset + 12 + length
if end > len(data):
raise ValueError("Truncated IMAGE PNG chunk")
tag = data[offset + 4:offset + 8]
body = data[offset + 8:end - 4]
if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"):
raise ValueError("Invalid IMAGE PNG checksum")
if tag == b"IHDR" and offset == 8 and length == 13:
width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body)
if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0):
raise ValueError("Unsupported IMAGE PNG layout")
dimensions = (height, width)
elif tag == b"IDAT" and dimensions:
saw_data = True
elif tag == b"IEND" and not length and end == len(data) and saw_data:
return dimensions
else:
raise ValueError("Unsupported IMAGE PNG chunk")
offset = end
raise ValueError("Truncated IMAGE PNG")
if fmt == "webp":
# Our encoder writes a single lossy VP8 still. Reject animation/alpha/
# extended WebP before WebPAnimDecoder can allocate its canvas.
if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ":
raise ValueError("Unsupported IMAGE WebP header")
size = int.from_bytes(data[16:20], "little")
if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data):
raise ValueError("Invalid IMAGE WebP length")
if data[23:26] != b"\x9d\x01\x2a":
raise ValueError("Invalid IMAGE WebP frame")
width, height = struct.unpack("<HH", data[26:30])
return height & 0x3fff, width & 0x3fff
if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
raise ValueError("Invalid IMAGE JPEG header or trailer")
offset, dimensions = 2, None
while offset + 4 <= len(data):
if data[offset] != 0xff:
raise ValueError("Invalid IMAGE JPEG marker")
marker = data[offset + 1]
size = int.from_bytes(data[offset + 2:offset + 4], "big")
end = offset + 2 + size
if size < 2 or end > len(data):
raise ValueError("Truncated IMAGE JPEG marker")
if marker in (0xc0, 0xc2):
if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3):
raise ValueError("Unsupported IMAGE JPEG layout")
dimensions = struct.unpack(">HH", data[offset + 5:offset + 9])
elif marker == 0xda:
if dimensions is None:
raise ValueError("Missing IMAGE JPEG dimensions")
return dimensions
offset = end
raise ValueError("Missing IMAGE JPEG pixels")
def _video_header(data, fmt, height, width, count):
"""Strict bounded EBML preflight: one track, no lacing/attachments/encodings.
Native stream probing is not the first dimension check: it may instantiate a
decoder. Validate the Matroska header and count blocks without decompressing.
"""
nodes = 0
def vint(pos, end, identifier=False):
if pos >= end or data[pos] == 0:
raise ValueError("Invalid IMAGE EBML integer")
length = 9 - data[pos].bit_length()
if length > (4 if identifier else 8) or pos + length > end:
raise ValueError("Truncated IMAGE EBML integer")
value = int.from_bytes(data[pos:pos + length], "big")
if not identifier:
value &= (1 << (7 * length)) - 1
if value == (1 << (7 * length)) - 1:
raise ValueError("Unknown-sized IMAGE EBML elements are forbidden")
return value, pos + length
def elements(start, end):
nonlocal nodes
while start < end:
nodes += 1
if nodes > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE container element count exceeds limit")
tag, pos = vint(start, end, True)
length, body = vint(pos, end)
finish = body + length
if finish > end:
raise ValueError("Truncated IMAGE EBML element")
yield tag, body, finish
start = finish
def integer(start, end):
if not 1 <= end - start <= 8:
raise ValueError("Invalid IMAGE EBML numeric field")
return int.from_bytes(data[start:end], "big")
roots = list(elements(0, len(data)))
if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]:
raise ValueError("Invalid IMAGE Matroska structure")
tracks, blocks, packets = [], 0, []
for tag, start, end in elements(roots[1][1], roots[1][2]):
if tag == 0x1654ae6b:
for entry, es, ee in elements(start, end):
if entry in (0xbf, 0xec):
continue
if entry != 0xae:
raise ValueError("Unexpected IMAGE track element")
fields, dims = {}, {}
for field, fs, fe in elements(es, ee):
if field in fields:
raise ValueError("Duplicate IMAGE track field")
fields[field] = (fs, fe)
if field == 0x6d80:
raise ValueError("IMAGE track encodings are forbidden")
if field == 0xe0:
for dim, ds, de in elements(fs, fe):
if dim in dims:
raise ValueError("Duplicate IMAGE video field")
dims[dim] = (ds, de)
if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)):
raise ValueError("Incomplete IMAGE video track")
if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1:
raise ValueError("Unsupported IMAGE video track")
if fields[0x86][1] - fields[0x86][0] > 32:
raise ValueError("IMAGE codec identifier exceeds limit")
codec_id = bytes(data[slice(*fields[0x86])])
if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"):
raise ValueError("IMAGE video codec mismatch")
if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height:
raise ValueError("IMAGE video header dimensions mismatch")
if codec_id == b"V_MS/VFW/FOURCC":
cs, ce = fields.get(0x63a2, (0, 0))
if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1":
raise ValueError("Invalid IMAGE FFV1 codec header")
if struct.unpack("<ii", data[cs + 4:cs + 12]) != (width, height):
raise ValueError("IMAGE FFV1 header dimensions mismatch")
tracks.append(fields)
elif tag == 0x1f43b675:
def block(bs, be):
nonlocal blocks
track, pos = vint(bs, be)
if track != 1 or pos + 3 >= be or data[pos + 2] & 0x06:
raise ValueError("Invalid or laced IMAGE video block")
blocks += 1
if blocks > count:
raise ValueError("IMAGE video frame count exceeds descriptor")
packets.append((pos + 3, be))
for field, fs, fe in elements(start, end):
if field == 0xa3:
block(fs, fe)
elif field == 0xa0:
group = [e for e in elements(fs, fe) if e[0] == 0xa1]
if len(group) != 1:
raise ValueError("Invalid IMAGE video block group")
block(group[0][1], group[0][2])
elif field not in (0xe7, 0xec, 0xbf):
raise ValueError("Unsupported IMAGE cluster field")
elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf):
raise ValueError("Unsupported IMAGE Matroska element")
if len(tracks) != 1 or blocks != count:
raise ValueError("IMAGE video track or frame count mismatch")
cs, ce = tracks[0].get(0x63a2, (0, 0))
if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES:
raise ValueError("IMAGE codec configuration exceeds limit or is missing")
if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC":
cs += 40 # Skip BITMAPINFOHEADER; FFV1 receives only its codec config.
return bytes(data[cs:ce]), packets
def _coded_frame_bound(data, fmt, header):
"""Conservative per-packet allocation bound before native decode.
FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one
VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU.
Counting slices overestimates pictures, which is safe. Do not trust a
Matroska block count alone: a malicious packet may contain several pictures.
"""
config, packets = header
if fmt == "ffv1":
return len(packets)
total, units = 0, 0
if fmt in ("h264", "h265"):
index = 4 if fmt == "h264" else 21
if len(config) <= index or config[0] != 1:
raise ValueError("Invalid IMAGE video configuration")
length_size = (config[index] & 3) + 1
for start, end in packets:
while start < end:
units += 1
if units > MAX_CONTAINER_ELEMENTS:
raise ValueError("IMAGE coded unit count exceeds limit")
if fmt in ("h264", "h265"):
if start + length_size > end:
raise ValueError("Truncated IMAGE NAL length")
length = int.from_bytes(data[start:start + length_size], "big")
start += length_size
if length < (1 if fmt == "h264" else 2) or length > end - start:
raise ValueError("Truncated IMAGE NAL unit")
kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63
total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31)
else:
tag = data[start]
start += 1
if tag & 0x81 or not tag & 2:
raise ValueError("Unsupported IMAGE AV1 OBU header")
kind = (tag >> 3) & 15
if tag & 4:
start += 1
length = 0
for shift in range(0, 56, 7):
if start >= end:
raise ValueError("Truncated IMAGE AV1 OBU length")
value = data[start]
start += 1
length |= (value & 127) << shift
if not value & 128:
break
else:
raise ValueError("IMAGE AV1 OBU length exceeds limit")
if length > end - start:
raise ValueError("Truncated IMAGE AV1 OBU")
total += kind in (3, 6)
start += length
if total > MAX_FRAMES * 4:
raise ValueError("IMAGE coded frame count exceeds limit")
if total < len(packets):
raise ValueError("IMAGE video has missing coded frames")
return total
def _decode_video(data, fmt, height, width, count, header):
av = _import_av()
maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64)
ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r")
ctx.width, ctx.height = width, height
ctx.extradata, packets = header
ctx.thread_count = 1
ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"}
ctx.open()
if "max_pixels" in ctx.options:
raise ValueError("IMAGE decoder does not support max_pixels safety limit")
if (ctx.width, ctx.height) != (width, height):
raise ValueError("IMAGE bitstream dimensions mismatch")
decoded = 0
for bounds in packets + [None]:
packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None
for frame in ctx.decode(packet):
if decoded >= count or (frame.width, frame.height) != (width, height):
raise ValueError("IMAGE decoded dimensions or frame count mismatch")
decoded += 1
yield frame.to_ndarray(format="rgb24")
if decoded != count:
raise ValueError("IMAGE decoded frame count mismatch")
def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
"""Decode to an independent contiguous CPU tensor after bounded preflight."""
limit = codec._byte_limit(max_bytes)
if not isinstance(descriptor, dict):
raise ValueError("IMAGE descriptor must be a dictionary")
if descriptor.get("encoding") == "tensor-le-v1":
return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit)
_descriptor_budget(descriptor)
if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}:
raise ValueError("Invalid IMAGE descriptor fields")
if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC":
raise ValueError("Invalid IMAGE descriptor encoding or layout")
fmt, _ = validate_options(descriptor["format"], descriptor["quality"])
if fmt == "exact":
raise ValueError("Exact IMAGE requires the tensor encoding")
if not isinstance(payload, bytes):
raise ValueError("IMAGE payload must be bytes")
if len(payload) > limit:
raise ValueError("IMAGE payload exceeds byte limit")
shape = descriptor["shape"]
if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES:
raise ValueError("Unsupported IMAGE descriptor dtype")
dtype, itemsize = codec._DTYPES[descriptor["dtype"]]
codec._tensor_size("IMAGE", shape, itemsize, limit)
alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit)
if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]:
raise ValueError("IMAGE alpha dtype mismatch")
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
count, height, width, channels = shape
lengths = descriptor["parts"]
if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count):
raise ValueError("IMAGE part count mismatch")
total = alpha_size
for length in lengths:
if type(length) is not int or not 0 < length <= limit - total:
raise ValueError("IMAGE part length exceeds byte limit")
total += length
if total != len(payload):
raise ValueError("IMAGE payload length does not match parts")
parts, offset = [], 0
for length in lengths:
parts.append(memoryview(payload)[offset:offset + length])
offset += length
try:
if fmt in _VIDEO_CODECS:
video_header = _video_header(parts[0], fmt, ph, pw, count)
coded_frames = _coded_frame_bound(parts[0], fmt, video_header)
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
else:
for part in parts:
if _still_header(part, fmt) != (height, width):
raise ValueError("IMAGE header dimensions mismatch")
# All headers and aggregate budgets are checked before output allocation.
alpha = None
if alpha_size:
alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit)
result = torch.empty(shape, dtype=dtype, device="cpu")
if fmt in _VIDEO_CODECS:
frames = _decode_video(parts[0], fmt, ph, pw, count, video_header)
else:
def stills():
import numpy as np
from PIL import Image
for part in parts:
with Image.open(io.BytesIO(part)) as image:
if image.size != (width, height) or getattr(image, "n_frames", 1) != 1:
raise ValueError("IMAGE still dimensions or frame count mismatch")
if image.mode not in ("L", "RGB"):
raise ValueError("Unsupported IMAGE still pixel mode")
yield np.array(image.convert("RGB"), dtype=np.uint8)
frames = stills()
color = 3 if channels == 4 else channels
for index, frame in enumerate(frames):
# Float32 division followed by the source dtype cast reproduces
# Collector normalized pixels without narrowing exact alpha.
result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255)
if alpha is not None:
result[..., 3:4] = alpha
return result
except ValueError:
raise
except Exception:
raise ValueError("Malformed or unsupported IMAGE compressed payload") from None
+30
View File
@@ -0,0 +1,30 @@
"""Image-return policy shared by the master and standalone endpoint bridge.
This module is deliberately dependency-free so graph compilation does not import
Torch, Pillow, PyAV or ComfyUI. Existing Collector settings stay unchanged.
"""
FORMATS = ("exact", "legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
DEFAULT_FORMAT = "exact"
DEFAULT_QUALITY = 90
FORMAT_TOOLTIP = (
"Image returns only; master-local images and other media stay unchanged. "
"exact: original floating-point tensor, no quantization. "
"legacy_png / png / ffv1: lossless 8-bit pixels, not exact float tensors. "
"jpeg / webp: lossy stills; h264 / h265 / av1: lossy frame-batch video. "
"Alpha is preserved separately when the codec cannot carry it. "
"Requires matching transfer support on the worker/endpoint."
)
QUALITY_TOOLTIP = (
"Quality 1–100 for lossy image returns; ignored by exact and lossless modes. "
"JPEG/WebP use it directly; video formats use the existing Collector CRF mapping."
)
def validate_options(transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
"""Reject malformed controls rather than repairing or silently substituting."""
if type(transfer_format) is not str or transfer_format not in FORMATS:
raise ValueError("Unsupported image transfer_format")
if type(transfer_quality) is not int or not 1 <= transfer_quality <= 100:
raise ValueError("Image transfer_quality must be an integer from 1 to 100")
return transfer_format, transfer_quality
+175
View File
@@ -0,0 +1,175 @@
"""Internal ComfyUI boundary nodes for the standalone media-v1 handler."""
import hashlib
import json
import uuid
from pathlib import Path
import folder_paths
from comfy_api.v0_0_2 import io
from .codec import decode_payload, encode_payload
from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
from .contract import (
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS,
bounded_int, identifier, validate_binding, validate_json,
)
def _scalar(value):
if isinstance(value, list):
if len(value) != 1:
raise ValueError("Media bridge control must have exactly one value")
return value[0]
return value
def _json(text):
if not isinstance(text, str) or len(text) > 256 * 1024:
raise ValueError("Invalid media bridge JSON")
try:
value = json.loads(text)
except (ValueError, TypeError, RecursionError):
raise ValueError("Invalid media bridge JSON") from None
validate_json(value)
return value
def _cpu(kind, value):
# Runtime transfer only, never sampling or device discovery. Pure codecs
# deliberately reject non-CPU storage rather than making hidden transfers.
if kind in ("IMAGE", "MASK"):
return value.detach().cpu()
if kind == "AUDIO":
return {"waveform": value["waveform"].detach().cpu(), "sample_rate": value["sample_rate"]}
return value
def encode_value(kind, value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
fmt, quality = validate_options(transfer_format, transfer_quality)
if kind == "IMAGE" and fmt != DEFAULT_FORMAT:
from .image_codec import encode_image
return encode_image(_cpu(kind, value), transfer_format=fmt, transfer_quality=quality)
if kind == "VIDEO":
from .video_codec import encode_video_cpu
return encode_video_cpu(value)
return encode_payload(kind, _cpu(kind, value))
def decode_value(kind, descriptor, payload):
if kind == "VIDEO":
from .video_codec import decode_video
return decode_video(descriptor, payload)
return decode_payload(kind, descriptor, payload)
def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
"""Write bounded owned payloads; history contains metadata, never base64."""
identifier(request_id)
validate_binding(binding)
fmt, quality = validate_options(transfer_format, transfer_quality)
if binding["kind"] != "IMAGE" and fmt != DEFAULT_FORMAT:
raise ValueError("Image transfer controls apply only to IMAGE exports")
if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS:
raise ValueError("Connected media output has no values or exceeds item limit")
request_key = hashlib.sha256(request_id.encode()).hexdigest()
subfolder = f"distributed-media-v1/{request_key}"
root = Path(folder_paths.get_temp_directory()).resolve()
directory = (root / subfolder).resolve()
if not directory.is_relative_to(root):
raise ValueError("Unsafe media output directory")
directory.mkdir(parents=True, exist_ok=True)
written, artifacts, items = [], [], []
total = 0
try:
for index, value in enumerate(values):
descriptor, payload = encode_value(binding["kind"], value,
transfer_format=fmt, transfer_quality=quality)
validate_json(descriptor)
total += len(payload)
if total > MAX_TOTAL_BYTES:
raise ValueError("Media output exceeds total byte limit")
aid = uuid.uuid4().hex
filename = aid + ".bin"
path = directory / filename
with path.open("xb") as stream:
written.append(path)
stream.write(payload)
artifacts.append({"id": aid, "kind": binding["kind"], "descriptor": descriptor,
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
"transport": {"type": "comfy-file", "filename": filename,
"subfolder": subfolder, "folder_type": "temp"}})
items.append({"binding_id": binding["binding_id"], "item_index": index, "artifact_id": aid})
except BaseException:
for path in written:
path.unlink(missing_ok=True)
raise
return {"request_id": request_id, "binding": dict(binding, count=len(values)),
"artifacts": artifacts, "items": items}
def import_value(artifact, filename):
"""Read a handler-staged, checksummed file within native media roots only."""
fields = {"id", "kind", "descriptor", "byte_length", "sha256"}
if not isinstance(artifact, dict) or set(artifact) != fields:
raise ValueError("Invalid staged media metadata")
identifier(artifact["id"])
if artifact["kind"] not in SOCKET_KINDS.values():
raise ValueError("Invalid staged media kind")
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
if not isinstance(filename, str):
raise ValueError("Invalid staged media filename")
roots = [Path(getter()).resolve() for getter in (
folder_paths.get_input_directory, folder_paths.get_output_directory, folder_paths.get_temp_directory)]
path = Path(folder_paths.get_annotated_filepath(filename)).resolve()
if not any(path.is_relative_to(root) for root in roots) or not path.is_file():
raise ValueError("Staged media file is outside allowed roots or unavailable")
if path.stat().st_size != size:
raise ValueError("Staged media file length mismatch")
with path.open("rb") as stream:
payload = stream.read(size + 1)
if len(payload) != size or hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
raise ValueError("Staged media file failed integrity validation")
return decode_value(artifact["kind"], artifact["descriptor"], payload)
class MediaImportV1(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
@classmethod
def execute(cls, artifact_json, filename):
artifact = _json(artifact_json)
value = import_value(artifact, filename)
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
class MediaExportV1(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="MediaExportV1", category="distributed/internal", is_deprecated=True,
inputs=[io.String.Input("request_id"), io.String.Input("binding_json"),
io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True),
io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True),
io.Video.Input("video", optional=True),
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=[], is_output_node=True, is_input_list=True)
@classmethod
def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None,
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
binding = _json(_scalar(binding_json))
validate_binding(binding)
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
if {name for name, value in provided.items() if value is not None} != {binding["socket"]}:
raise ValueError("Media export must receive exactly its declared typed socket")
exported = export_values(_scalar(request_id), binding, provided[binding["socket"]],
transfer_format=_scalar(transfer_format),
transfer_quality=_scalar(transfer_quality))
return io.NodeOutput(ui={"media_exports": [exported]})
+482
View File
@@ -0,0 +1,482 @@
"""Lossless transport of pinned native ComfyUI VIDEO representations.
No save_to/get_components fallback is allowed: an unknown class must not silently
transcode. Tensor components use the shared CPU codec. File bytes are opaque and
retain signed trim/crop parameters; decoded files own independent BytesIOs.
VideoFromList's native constructor flattens its inputs, so wire segments are a
bounded, flat list of exact VideoFromFile objects, never recursive playlists.
Offsets are relative to the containing VIDEO payload. Canonical part order is
images/alpha/audio, or segments/complete_audio. Every byte is covered exactly
once. All descriptor structure and tensor sizes are preflighted before decoding
any media storage. File containers are not parsed until a native consumer reads
them; this adapter does not attempt to validate or repair container codecs.
"""
from fractions import Fraction
import io
import json
import math
import os
from pathlib import Path
import stat
from .codec import MAX_PAYLOAD_BYTES, decode_payload, encode_payload
MAX_SEGMENTS = 64
MAX_METADATA_BYTES = 64 * 1024
MAX_METADATA_DEPTH = 16
_ENCODING = "native-video-v1"
_BASE_FIELDS = {"encoding", "representation", "byte_length"}
_COMPONENT_FIELDS = {"images", "alpha", "audio", "frame_rate", "metadata", "bit_depth", "color_space"}
_FILE_FIELDS = {"data", "start_time", "duration", "crop"}
_SEGMENT_FIELDS = {"segments", "complete_audio"}
_TENSOR_BYTES = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}
def _limit(max_bytes):
if type(max_bytes) is not int or max_bytes < 0:
raise ValueError("VIDEO max_bytes must be a nonnegative integer")
return min(max_bytes, MAX_PAYLOAD_BYTES)
def _native_types():
# Lazy imports keep non-VIDEO users independent of ComfyUI and PyAV.
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
return VideoFromComponents, VideoFromFile, VideoFromList, VideoComponents
def _fields(value, fields, label):
if type(value) is not dict or set(value) != fields:
raise ValueError(f"Invalid VIDEO {label} fields")
def _integer(value, *, positive=False):
if type(value) is not int or value < (1 if positive else 0):
raise ValueError("VIDEO length, offset and dimensions must be integers in range")
return value
def _number(value):
if type(value) not in (int, float) or (type(value) is float and not math.isfinite(value)):
raise ValueError("VIDEO trim values must be finite numbers")
return value
def _crop(value, *, wire=False):
if value is None:
return None
if type(value) is not (list if wire else tuple) or len(value) != 4:
raise ValueError("VIDEO crop must contain four integers")
if any(type(part) is not int for part in value):
raise ValueError("VIDEO crop must contain four integers")
# Native crop normalization is a consumer concern, not a transport operation.
return list(value) if not wire else tuple(value)
def _frame_rate(value):
_fields(value, {"numerator", "denominator"}, "frame_rate")
_integer(value["numerator"], positive=True)
_integer(value["denominator"], positive=True)
def _encoding_options(bit_depth, color_space):
_integer(bit_depth, positive=True)
if type(color_space) is not str or color_space not in ("sRGB", "HDR", "HDR PQ"):
raise ValueError("Unsupported native VIDEO color_space")
def _metadata(value):
"""Validate bounded JSON before allocating its serialized form; return a copy."""
if value is None:
return None
if type(value) is not dict:
raise ValueError("VIDEO metadata must be a JSON object or null")
remaining = MAX_METADATA_BYTES
def charge(size):
nonlocal remaining
remaining -= size
if remaining < 0:
raise ValueError("VIDEO metadata exceeds byte limit")
def visit(item, depth):
if depth > MAX_METADATA_DEPTH:
raise ValueError("VIDEO metadata exceeds depth limit")
kind = type(item)
if kind in (dict, list):
# Each member costs at least one JSON byte. Also bounds iteration
# before walking potentially huge containers or cyclic structures.
if len(item) > remaining:
raise ValueError("VIDEO metadata exceeds byte limit")
charge(2 + max(0, len(item) - 1))
if kind is dict:
for key, child in item.items():
if type(key) is not str:
raise ValueError("VIDEO metadata keys must be strings")
visit(key, depth + 1)
charge(1)
visit(child, depth + 1)
else:
for child in item:
visit(child, depth + 1)
elif kind in (str, int, float, bool) or item is None:
if kind is str and len(item) > remaining:
raise ValueError("VIDEO metadata exceeds byte limit")
# Avoid converting an unbounded Python integer to decimal first.
if kind is int and item.bit_length() > MAX_METADATA_BYTES * 4:
raise ValueError("VIDEO metadata exceeds byte limit")
if kind is float and not math.isfinite(item):
raise ValueError("VIDEO metadata numbers must be finite")
try:
charge(len(json.dumps(item, ensure_ascii=True, allow_nan=False)))
except (ValueError, OverflowError):
raise ValueError("Invalid or oversized VIDEO metadata scalar") from None
else:
raise ValueError("VIDEO metadata must contain only JSON values")
visit(value, 0)
return json.loads(json.dumps(value, ensure_ascii=True, allow_nan=False, separators=(",", ":")))
def _attrs(value, names):
attributes = vars(value)
if not all(name in attributes for name in names):
raise ValueError("Native VIDEO attributes do not match the pinned adapter")
return [attributes[name] for name in names]
def _roots(allowed_roots):
if allowed_roots is None:
import folder_paths
allowed_roots = [folder_paths.get_input_directory(), folder_paths.get_output_directory(),
folder_paths.get_temp_directory()]
if type(allowed_roots) not in (list, tuple):
raise ValueError("VIDEO allowed_roots must be a list or tuple")
try:
roots = [Path(root).resolve() for root in allowed_roots]
except (OSError, TypeError, ValueError, RuntimeError):
raise ValueError("Invalid VIDEO allowed roots") from None
return roots
def _read_path(value, limit, allowed_roots):
"""Resolve authorization, reject special files, and bound all reads.
Open each resolved path component with NOFOLLOW using directory fds, so an
intermediate symlink swapped in after resolve() cannot bypass authorization.
NONBLOCK prevents an attacker swapping in a FIFO from hanging the worker.
"""
roots = _roots(allowed_roots)
directory_fd = file_fd = None
try:
path = Path(value).resolve(strict=True)
if not any(path.is_relative_to(root) for root in roots):
raise ValueError("VIDEO file is outside allowed roots")
directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
directory_fd = os.open(path.anchor, directory_flags)
for part in path.parts[1:-1]:
next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
os.close(directory_fd)
directory_fd = next_fd
file_fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd)
info = os.fstat(file_fd)
if not stat.S_ISREG(info.st_mode):
raise ValueError("VIDEO source must be a regular file")
if info.st_size > limit:
raise ValueError("VIDEO file exceeds byte limit")
# Read at most the stat size plus a sentinel; reject growth or truncation.
chunks = []
count = 0
while count <= info.st_size:
chunk = os.read(file_fd, min(1024 * 1024, info.st_size + 1 - count))
if not chunk:
break
chunks.append(chunk)
count += len(chunk)
if count != info.st_size:
raise ValueError("VIDEO source changed size during bounded read")
return b"".join(chunks)
except (OSError, RuntimeError):
raise ValueError("VIDEO file cannot be read safely") from None
finally:
if file_fd is not None:
os.close(file_fd)
if directory_fd is not None:
os.close(directory_fd)
def _file_bytes(value, limit, allowed_roots):
if type(value) is io.BytesIO:
try:
with value.getbuffer() as view:
if view.nbytes > limit:
raise ValueError("VIDEO file exceeds byte limit")
return bytes(view)
except ValueError:
raise ValueError("VIDEO BytesIO is closed or exceeds byte limit") from None
if isinstance(value, (str, os.PathLike)):
return _read_path(value, limit, allowed_roots)
raise ValueError("VIDEO file must be a local path or exact BytesIO")
def encode_video(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
"""Encode an exact native VIDEO without opening/transcoding its container.
Paths are authorized only under allowed_roots (when supplied), otherwise
under ComfyUI input/output/temp roots. Budgets may lower, never raise, the
shared maximum. Unsupported native state raises ValueError, not a fallback.
Non-CPU tensors are rejected; runtime staging belongs to the bridge boundary.
"""
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=False)
def encode_video_cpu(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
"""Bridge-only boundary: detach/stage native tensors, then losslessly encode.
Share the exact-type adapter with encode_video rather than reconstructing
native videos, calling get_components, or duplicating private attributes.
Images, alpha and both audio representations move without dtype conversion;
metadata/options and opaque file bytes/trim/crop take the unchanged path.
"""
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=True)
def _stage_tensor_cpu(value):
import torch
if not isinstance(value, torch.Tensor):
raise ValueError("Tensor value must be a torch.Tensor")
return value.detach().cpu()
def _encode_video(value, *, max_bytes, allowed_roots, stage_cpu):
limit = _limit(max_bytes)
components_class, file_class, list_class, components_type = _native_types()
descriptor = {"encoding": _ENCODING}
chunks = []
size = 0
def add(payload, nested=None):
nonlocal size
if len(payload) > limit - size:
raise ValueError("VIDEO payload exceeds byte limit")
slot = {"offset": size, "length": len(payload)}
if nested is not None:
slot["descriptor"] = nested
chunks.append(payload)
size += len(payload)
return slot
def tensor(kind, tensor_value):
if stage_cpu:
if kind == "AUDIO":
if not isinstance(tensor_value, dict) or set(tensor_value) != {"waveform", "sample_rate"}:
raise ValueError("AUDIO value must contain waveform and sample_rate")
tensor_value = {**tensor_value, "waveform": _stage_tensor_cpu(tensor_value["waveform"])}
else:
tensor_value = _stage_tensor_cpu(tensor_value)
nested, payload = encode_payload(kind, tensor_value, max_bytes=limit - size)
return add(payload, nested)
if type(value) is components_class:
component, bit_depth, color_space = _attrs(value, [
"_VideoFromComponents__components", "_VideoFromComponents__bit_depth",
"_VideoFromComponents__color_space",
])
if type(component) is not components_type:
raise ValueError("Unsupported native VIDEO components class")
images, fps, audio, metadata, alpha = _attrs(component, [
"images", "frame_rate", "audio", "metadata", "alpha",
])
if type(fps) is not Fraction or fps <= 0:
raise ValueError("VIDEO frame_rate must be a positive Fraction")
_encoding_options(bit_depth, color_space)
descriptor.update(representation="components", bit_depth=bit_depth, color_space=color_space,
frame_rate={"numerator": fps.numerator, "denominator": fps.denominator},
metadata=_metadata(metadata))
descriptor["images"] = tensor("IMAGE", images)
# Native get_components() returns BHW1 alpha; MASK nodes use BHW.
# Select an existing lossless tensor layout without squeezing/converting.
alpha_kind = "IMAGE" if getattr(alpha, "ndim", None) == 4 else "MASK"
descriptor["alpha"] = tensor(alpha_kind, alpha) if alpha is not None else None
if alpha is not None:
_alpha_header(descriptor["alpha"], descriptor["images"])
descriptor["audio"] = tensor("AUDIO", audio) if audio is not None else None
elif type(value) is file_class:
source, start, duration, crop = _attrs(value, [
"_VideoFromFile__file", "_VideoFromFile__start_time", "_VideoFromFile__duration",
"_VideoFromFile__crop",
])
descriptor.update(representation="file", start_time=_number(start), duration=_number(duration),
crop=_crop(crop))
descriptor["data"] = add(_file_bytes(source, limit, allowed_roots))
elif type(value) is list_class:
videos, complete_audio = _attrs(value, ["videos", "complete_audio"])
if type(videos) is not list or not 1 <= len(videos) <= MAX_SEGMENTS:
raise ValueError("VIDEO must have between 1 and 64 segments")
if any(type(segment) is not file_class for segment in videos):
raise ValueError("VIDEO segments must be exact native VideoFromFile objects")
descriptor.update(representation="segments", segments=[])
for segment in videos:
nested, payload = encode_video(segment, max_bytes=limit - size, allowed_roots=allowed_roots)
descriptor["segments"].append(add(payload, nested))
descriptor["complete_audio"] = tensor("AUDIO", complete_audio) if complete_audio is not None else None
else:
raise ValueError("Unsupported native VIDEO class")
descriptor["byte_length"] = size
return descriptor, b"".join(chunks)
def _tensor_header(kind, descriptor, length):
"""Allocation-free preflight of the shared tensor-le-v1 descriptor.
The shared decoder still validates the bytes/finite values. Doing its small
structural check here first prevents allocation for an early part when a
later part's descriptor is already known to be invalid.
"""
fields = {"encoding", "dtype", "shape", "layout"}
if kind == "AUDIO":
fields.add("sample_rate")
_fields(descriptor, fields, "tensor")
layout = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}[kind]
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != "tensor-le-v1":
raise ValueError("Invalid VIDEO tensor encoding")
if type(descriptor["layout"]) is not str or descriptor["layout"] != layout:
raise ValueError("Invalid VIDEO tensor layout")
dtype = descriptor["dtype"]
if type(dtype) is not str or dtype not in _TENSOR_BYTES:
raise ValueError("Invalid VIDEO tensor dtype")
shape = descriptor["shape"]
if type(shape) is not list or len(shape) != len(layout):
raise ValueError("Invalid VIDEO tensor rank")
size = _TENSOR_BYTES[dtype]
for dimension in shape:
_integer(dimension, positive=True)
if dimension > length // size:
raise ValueError("VIDEO tensor shape exceeds byte limit")
size *= dimension
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
raise ValueError("Invalid VIDEO image channels")
if size != length:
raise ValueError("VIDEO tensor length does not match shape")
if kind == "AUDIO":
_integer(descriptor["sample_rate"], positive=True)
def _alpha_header(alpha, images):
_fields(alpha, {"offset", "length", "descriptor"}, "alpha part")
descriptor = alpha["descriptor"]
if type(descriptor) is not dict:
raise ValueError("Invalid VIDEO alpha descriptor")
layout = descriptor.get("layout")
if type(layout) is not str or layout not in ("BHW", "BHWC"):
raise ValueError("Invalid VIDEO alpha layout")
kind = "MASK" if layout == "BHW" else "IMAGE"
_tensor_header(kind, descriptor, _integer(alpha["length"]))
expected = images["descriptor"]["shape"][:3] + ([1] if kind == "IMAGE" else [])
if descriptor["shape"] != expected:
raise ValueError("VIDEO alpha shape must match images")
return kind
def _validate_video(descriptor, size, *, file_only=False):
if type(descriptor) is not dict:
raise ValueError("VIDEO descriptor must be a dictionary")
representation = descriptor.get("representation")
if type(representation) is not str or representation not in ("components", "file", "segments"):
raise ValueError("Unsupported VIDEO representation")
if file_only and representation != "file":
# Depth is at most one: native lists already flatten their segments.
raise ValueError("VIDEO segments cannot contain nested lists or components")
fields = {"components": _COMPONENT_FIELDS, "file": _FILE_FIELDS, "segments": _SEGMENT_FIELDS}[representation]
_fields(descriptor, _BASE_FIELDS | fields, "descriptor")
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != _ENCODING:
raise ValueError("Unsupported VIDEO encoding")
if _integer(descriptor["byte_length"]) != size:
raise ValueError("VIDEO payload length does not match descriptor")
result = dict(descriptor)
cursor = 0
def part(slot, kind=None, video=False):
nonlocal cursor
fields = {"offset", "length", "descriptor"} if kind or video else {"offset", "length"}
_fields(slot, fields, "part")
offset = _integer(slot["offset"])
length = _integer(slot["length"])
if offset != cursor or length > size - cursor:
raise ValueError("VIDEO parts must not overlap, leave gaps or exceed bounds")
cursor += length
if kind:
_tensor_header(kind, slot["descriptor"], length)
if video:
return {**slot, "descriptor": _validate_video(slot["descriptor"], length, file_only=True)}
return slot
if representation == "components":
_frame_rate(descriptor["frame_rate"])
_encoding_options(descriptor["bit_depth"], descriptor["color_space"])
result["metadata"] = _metadata(descriptor["metadata"])
part(descriptor["images"], "IMAGE")
if descriptor["alpha"] is not None:
part(descriptor["alpha"], _alpha_header(descriptor["alpha"], descriptor["images"]))
if descriptor["audio"] is not None:
part(descriptor["audio"], "AUDIO")
elif representation == "file":
_number(descriptor["start_time"])
_number(descriptor["duration"])
_crop(descriptor["crop"], wire=True)
part(descriptor["data"])
else:
segments = descriptor["segments"]
if type(segments) is not list or not 1 <= len(segments) <= MAX_SEGMENTS:
raise ValueError("VIDEO must have between 1 and 64 segments")
result["segments"] = [part(segment, video=True) for segment in segments]
if descriptor["complete_audio"] is not None:
part(descriptor["complete_audio"], "AUDIO")
if cursor != size:
raise ValueError("VIDEO payload contains trailing bytes")
return result
def _decode_validated(descriptor, payload, types):
components_class, file_class, list_class, components_type = types
def data(slot):
return payload[slot["offset"]:slot["offset"] + slot["length"]]
def tensor(kind, slot):
if slot is None:
return None
return decode_payload(kind, slot["descriptor"], data(slot), max_bytes=slot["length"])
if descriptor["representation"] == "file":
# VideoFromFile retains the BytesIO, unlike a temporary filesystem path.
return file_class(io.BytesIO(data(descriptor["data"])), start_time=descriptor["start_time"],
duration=descriptor["duration"], crop=_crop(descriptor["crop"], wire=True))
if descriptor["representation"] == "components":
fps = descriptor["frame_rate"]
return components_class(components_type(
images=tensor("IMAGE", descriptor["images"]),
frame_rate=Fraction(fps["numerator"], fps["denominator"]),
audio=tensor("AUDIO", descriptor["audio"]),
metadata=descriptor["metadata"],
alpha=tensor("IMAGE" if descriptor["alpha"] is not None and
descriptor["alpha"]["descriptor"]["layout"] == "BHWC" else "MASK",
descriptor["alpha"]),
), bit_depth=descriptor["bit_depth"], color_space=descriptor["color_space"])
videos = [_decode_validated(slot["descriptor"], data(slot), types) for slot in descriptor["segments"]]
return list_class(videos, complete_audio=tensor("AUDIO", descriptor["complete_audio"]))
def decode_video(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
"""Decode a bounded VIDEO into independently owned native CPU media objects."""
limit = _limit(max_bytes)
if type(payload) is not bytes:
raise ValueError("VIDEO payload must be bytes")
if len(payload) > limit:
raise ValueError("VIDEO payload exceeds byte limit")
validated = _validate_video(descriptor, len(payload))
return _decode_validated(validated, payload, _native_types())
+10 -11
View File
@@ -38,17 +38,10 @@ def _load_config_routes_module():
sys.modules[f"{package_name}.api"] = api_pkg
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
utils_pkg.__path__ = [str(module_path.parents[1] / "utils")]
sys.modules[f"{package_name}.utils"] = utils_pkg
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
import aiohttp # Load the real dependency; isolate the fake response at the tested module.
class _Routes:
def get(self, _path):
@@ -96,8 +89,14 @@ def _load_config_routes_module():
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)
class _EmptySecretStore:
def has(self, _worker_id):
return False
module.web = types.SimpleNamespace(
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
)
module._serverless_secret_store = lambda: _EmptySecretStore()
return module
+233 -3
View File
@@ -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
@@ -162,7 +176,7 @@ def _load_job_routes_module():
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", 1))
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)
@@ -205,6 +219,10 @@ def _load_job_routes_module():
module = importlib.util.module_from_spec(spec)
assert spec is not None and spec.loader is not None
spec.loader.exec_module(module)
# Keep response assertions independent of real aiohttp import ordering.
module.web = types.SimpleNamespace(
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
)
if created_aiohttp_stub:
sys.modules.pop("aiohttp", None)
@@ -221,7 +239,7 @@ job_routes = _load_job_routes_module()
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
async def test_distributed_queue_happy_path_returns_prompt_id(self):
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
request = _FakeRequest(
{
"prompt": {"1": {"class_type": "Node"}},
@@ -233,12 +251,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
with patch.object(
job_routes,
"orchestrate_distributed_execution",
new=AsyncMock(return_value=("prompt_123", 2)),
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):
@@ -300,6 +320,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
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,
@@ -321,5 +412,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()
+116 -15
View File
@@ -164,26 +164,11 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
# ---------------------------------------------------------------------------
class FindMediaReferencesTests(unittest.TestCase):
def test_finds_image_input(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
refs = ms._find_media_references(prompt)
self.assertIn("photo.png", refs)
def test_finds_video_input(self):
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
refs = ms._find_media_references(prompt)
self.assertIn("clip.mp4", refs)
def test_finds_file_input_for_load_video(self):
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
refs = ms._find_media_references(prompt)
self.assertIn("1 - Copy.mp4", refs)
def test_finds_audio_input(self):
prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
refs = ms._find_media_references(prompt)
self.assertIn("track.wav", refs)
def test_strips_annotation_suffix(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
refs = ms._find_media_references(prompt)
@@ -258,5 +243,121 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase):
async def test_output_temp_and_input_uploads_do_not_collide(self):
from unittest.mock import AsyncMock, patch
class Response:
status = 200
def __init__(self, data):
self.data = data
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.data
def raise_for_status(self):
pass
class Form:
def __init__(self):
self.fields = {}
def add_field(self, key, value, **kwargs):
self.fields[key] = value
if key == 'image':
self.fields['name'] = kwargs['filename']
checked = []
def post(url, *, json=None, data=None, **kwargs):
if url.endswith('/check_file'):
checked.append(json['filename'])
return Response({'exists': False})
self.assertEqual(data.fields['type'], 'input')
self.assertEqual(data.fields['image'], b'real file bytes')
return Response({'name': data.fields['name'], 'subfolder': data.fields['subfolder']})
refs = ['shared/ref.png', 'shared/ref.png [output]', 'shared/ref.png [temp]']
prompt = {str(i): {'inputs': {'image': ref}} for i, ref in enumerate(refs)}
self.assertEqual(ms._find_media_references(prompt), refs)
with patch.object(ms, 'get_client_session', AsyncMock(return_value=types.SimpleNamespace(post=post))), \
patch.object(ms, 'aiohttp', types.SimpleNamespace(FormData=Form, ClientTimeout=lambda **kw: kw)):
paths = [await ms._upload_media_to_worker({}, ref, b'real file bytes', 'fixture-hash', 'image/png')
for ref in refs]
expected = ['shared/ref.png', 'distributed-media/output/fixture-hash/ref.png',
'distributed-media/temp/fixture-hash/ref.png']
self.assertEqual(checked, expected)
self.assertEqual(paths, [(True, path) for path in expected])
ms._rewrite_prompt_media_inputs(prompt, dict(zip(refs, expected)))
self.assertEqual([node['inputs']['image'] for node in prompt.values()], expected)
if __name__ == "__main__":
unittest.main()
# ---------------------------------------------------------------------------
# fetch_worker_has_node
# ---------------------------------------------------------------------------
class _FakeResponse:
def __init__(self, status, payload):
self.status = status
self._payload = payload
async def json(self):
if isinstance(self._payload, Exception):
raise self._payload
return self._payload
async def __aenter__(self):
return self
async def __aexit__(self, *exc):
return False
class _FakeSession:
def __init__(self, response=None, error=None):
self.response = response
self.error = error
self.urls = []
def get(self, url, timeout=None):
self.urls.append(url)
if self.error:
raise self.error
return self.response
class FetchWorkerHasNodeTests(unittest.TestCase):
def _probe(self, session):
import asyncio
async def _get_session():
return session
original = ms.get_client_session
ms.get_client_session = _get_session
try:
return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
finally:
ms.get_client_session = original
def test_registered_node_is_reported(self):
session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
self.assertTrue(self._probe(session))
self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
def test_unknown_node_returns_empty_object(self):
# ComfyUI answers 200 {} for classes that are not registered.
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
def test_http_error_counts_as_missing(self):
self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
def test_bad_json_counts_as_missing(self):
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
def test_connection_error_counts_as_missing(self):
self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))
+219
View File
@@ -0,0 +1,219 @@
"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
import asyncio
import importlib.util
import sys
from pathlib import Path
import numpy as np
import pytest
import torch
_TESTS_DIR = Path(__file__).resolve().parents[1]
def _load_from(path, name):
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _load_modules():
real_aiohttp = sys.modules["aiohttp"] # imported inside the isolated test
queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
job_routes = queue_tests._load_job_routes_module()
collector = collector_tests._load_collector_module()
# Error responses must be real HTTP responses here (the unit-test harness
# returns plain objects); mirror utils.network.handle_api_error.
from aiohttp import web
async def handle_api_error(_request, error, status=500):
if isinstance(error, list):
return web.json_response({"errors": [str(item) for item in error]}, status=status)
return web.json_response({"error": str(error)}, status=status)
job_routes.handle_api_error = handle_api_error
# The queue-test harness also fakes aiohttp.web for route modules; the
# endpoint must build real responses for the real server used here.
job_routes.web = web
# The collector harness stubs aiohttp; this test needs the real client.
sys.modules["aiohttp"] = real_aiohttp
collector.aiohttp = real_aiohttp
return job_routes, collector
@pytest.fixture(autouse=True)
def _isolate_sys_modules():
"""The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
saved = dict(sys.modules)
# Earlier test modules may have left a stub `aiohttp` behind; drop it so the
# real package is imported for this test (restored afterwards).
for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
if getattr(sys.modules[name], "__file__", None) is None:
del sys.modules[name]
yield
sys.modules.clear()
sys.modules.update(saved)
async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
request_limit, chunk_target, audio=None):
"""Send `images` from a real worker node to the real completion route.
Returns (queued items, sizes of every request body the server received,
exception raised by the worker or None).
"""
import aiohttp
from aiohttp import web
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
body_sizes = []
@web.middleware
async def record_size(request, handler):
body_sizes.append(request.content_length)
return await handler(request)
app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
error = None
try:
async with aiohttp.ClientSession() as session:
async def _get_session():
return session
collector_module.get_client_session = _get_session
collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
node = collector_module.DistributedCollectorNode()
try:
await node.send_batch_to_master(
images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
)
except Exception as exc: # surfaced to the caller for assertions
error = exc
finally:
await runner.cleanup()
items = []
while not queue.empty():
items.append(queue.get_nowait())
return items, body_sizes, error
def _noise_images(count, height=360, width=640, seed=97):
return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
def test_high_entropy_video_stays_under_server_upload_limit(fmt):
"""Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(24)
limit = 2 * 1024 * 1024
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, fmt,
client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
assert error is None, error
assert len(body_sizes) > 1 and max(body_sizes) <= limit
assert [item["image_index"] for item in items] == list(range(24))
assert [item["is_last"] for item in items] == [False] * 23 + [True]
def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
"""A worker configured above the master's limit fails loudly instead of
silently resending the whole batch as uncompressed PNG."""
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(8)
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, "ffv1",
client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
assert isinstance(error, RuntimeError)
assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
assert len(body_sizes) == 1
assert items == []
def test_frame_that_cannot_fit_sends_nothing():
pytest.importorskip("aiohttp")
job_routes, collector_module = _load_modules()
images = _noise_images(3)
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
items, body_sizes, error = asyncio.run(_run_transfer(
job_routes, collector_module, images, "png",
client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
assert body_sizes == []
assert items == []
@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
def test_worker_to_master_over_real_http(fmt):
# Import here, not at module level: other test harnesses only install their
# aiohttp stubs when the real package has not been imported yet.
aiohttp = pytest.importorskip("aiohttp")
from aiohttp import web
job_routes, collector_module = _load_modules()
images = torch.rand(5, 24, 36, 3)
async def scenario():
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
app = web.Application()
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
async with aiohttp.ClientSession() as session:
async def _get_session():
return session
collector_module.get_client_session = _get_session
collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000 # several requests for still formats
node = collector_module.DistributedCollectorNode()
await node.send_batch_to_master(
images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
)
await runner.cleanup()
items = []
while not queue.empty():
items.append(queue.get_nowait())
return items
items = asyncio.run(scenario())
assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
assert [item["is_last"] for item in items] == [False, False, False, False, True]
expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
got = torch.cat([item["tensor"] for item in items]).numpy()
assert got.shape == expected.shape
if fmt in ("png", "ffv1"):
assert np.array_equal(got, expected)
else:
assert np.abs(got - expected).mean() < 0.25 # random noise: only sanity-check lossy codecs
+104 -1
View File
@@ -106,7 +106,7 @@ def _load_worker_routes_module():
sys.modules[f"{package_name}.utils"] = utils_pkg
workers_pkg = types.ModuleType(f"{package_name}.workers")
workers_pkg.__path__ = []
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
sys.modules[f"{package_name}.workers"] = workers_pkg
@@ -213,6 +213,7 @@ def _load_worker_routes_module():
raise RuntimeError("not used in these tests")
network_module.get_client_session = _get_client_session
network_module.get_server_port = lambda: 8189
sys.modules[f"{package_name}.utils.network"] = network_module
constants_module = types.ModuleType(f"{package_name}.utils.constants")
@@ -256,6 +257,8 @@ def _load_worker_routes_module():
module = importlib.util.module_from_spec(spec)
assert spec is not None and spec.loader is not None
spec.loader.exec_module(module)
# Keep this module's response double independent of aiohttp import order.
module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
if created_aiohttp_stub:
sys.modules.pop("aiohttp", None)
@@ -268,7 +271,107 @@ def _load_worker_routes_module():
worker_routes = _load_worker_routes_module()
class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
from test_async_helpers import _load_async_helpers_module
self.modules = patch.dict(sys.modules)
self.modules.start()
self.addCleanup(self.modules.stop)
self.helper, self.server = _load_async_helpers_module()
sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper
self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}}
async def receive(self, payload):
import json
sent = []
class Socket:
async def prepare(self, request):
pass
async def send_json(self, value):
sent.append(value)
def __aiter__(self):
async def messages():
yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT,
data=json.dumps(payload))
return messages()
with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \
patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \
patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError):
await worker_routes.worker_ws_endpoint(_FakeRequest())
return sent[0]
def payload(self, targets):
return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt,
'partial_execution_targets': targets}
async def test_explicit_targets_reach_queue_and_acknowledgement(self):
ack = await self.receive(self.payload(['9']))
self.assertTrue(ack['ok'], ack)
self.assertEqual(ack['request_id'], 'req')
self.assertEqual(self.server.prompt_queue.items[0][4], ['9'])
async def test_malformed_targets_and_messages_are_rejected(self):
for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']):
with self.subTest(targets=targets):
ack = await self.receive(self.payload(targets))
self.assertFalse(ack['ok'])
self.assertIn('partial_execution_targets', ack['error'])
self.assertFalse((await self.receive([]))['ok'])
self.assertEqual(self.server.prompt_queue.items, [])
async def test_validation_cannot_silently_queue_a_partial_subset(self):
ack = await self.receive(self.payload(['9', '10']))
self.assertFalse(ack['ok'])
self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation')
self.assertEqual(self.server.prompt_queue.items, [])
async def test_missing_target_field_keeps_legacy_queue_call(self):
from unittest.mock import AsyncMock
queued = AsyncMock(return_value='legacy')
original = self.helper.queue_prompt_payload
self.helper.queue_prompt_payload = queued
try:
payload = self.payload(['9'])
payload.pop('partial_execution_targets')
self.assertTrue((await self.receive(payload))['ok'])
self.assertNotIn('partial_execution_targets', queued.call_args.kwargs)
finally:
self.helper.queue_prompt_payload = original
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
async def test_network_info_reports_actual_master_and_available_ports(self):
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
patch.object(worker_routes, "load_config", return_value={"workers": []}), \
patch.object(worker_routes, "get_server_port", return_value=8189), \
patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
self.assertEqual(response.payload["master_port"], 8189)
self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
self.assertEqual(response.payload["cuda_device_count"], 4)
allocate.assert_called_once_with(8189, [], 3)
async def test_network_info_does_not_reallocate_existing_configuration(self):
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
patch.object(worker_routes, "get_network_ips", return_value=[]), \
patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
patch.object(worker_routes, "allocate_worker_ports") as allocate:
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
self.assertEqual(response.payload["local_worker_ports"], [])
allocate.assert_not_called()
async def test_launch_conflict_is_a_clear_client_error(self):
manager = _DummyWorkerManager()
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
patch.object(worker_routes, "load_config", return_value=config), \
patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
self.assertEqual(response.status, 400)
self.assertIn("master port 8189", response.payload["message"])
self.assertEqual(manager.processes, {})
async def test_launch_worker_valid_id_returns_200(self):
manager = _DummyWorkerManager()
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
+691
View File
@@ -0,0 +1,691 @@
{
"DistributedCollector": {
"input": {
"required": {
"load_balance": [
"BOOLEAN",
{
"default": false,
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
}
]
},
"optional": {
"images": [
"IMAGE"
],
"audio": [
"AUDIO"
]
},
"hidden": {
"multi_job_id": [
"STRING",
{
"default": ""
}
],
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"master_url": [
"STRING",
{
"default": ""
}
],
"enabled_worker_ids": [
"STRING",
{
"default": "[]"
}
],
"worker_batch_size": [
"INT",
{
"default": 1,
"min": 1,
"max": 1024
}
],
"worker_id": [
"STRING",
{
"default": ""
}
],
"pass_through": [
"BOOLEAN",
{
"default": false
}
],
"delegate_only": [
"BOOLEAN",
{
"default": false
}
]
}
},
"input_order": {
"required": [
"load_balance"
],
"optional": [
"images",
"audio"
],
"hidden": [
"multi_job_id",
"is_worker",
"master_url",
"enabled_worker_ids",
"worker_batch_size",
"worker_id",
"pass_through",
"delegate_only"
]
},
"output": [
"IMAGE",
"AUDIO"
],
"output_name": [
"images",
"audio"
],
"output_is_list": [
false,
false
],
"is_input_list": true,
"output_node": false,
"category": "image",
"display_name": "Distributed Collector"
},
"DistributedSeed": {
"input": {
"required": {
"seed": [
"INT",
{
"default": 1125899906842,
"min": 0,
"max": 1125899906842624,
"forceInput": false
}
]
},
"hidden": {
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"worker_id": [
"STRING",
{
"default": ""
}
]
}
},
"input_order": {
"required": [
"seed"
],
"hidden": [
"is_worker",
"worker_id"
]
},
"output": [
"INT"
],
"output_name": [
"seed"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "utils",
"display_name": "Distributed Seed"
},
"DistributedModelName": {
"input": {
"required": {
"text": [
"STRING",
{
"default": ""
}
]
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO"
}
},
"input_order": {
"required": [
"text"
],
"hidden": [
"unique_id",
"extra_pnginfo"
]
},
"output": [
"*"
],
"output_name": [
"output"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": true,
"category": "utils",
"display_name": "Distributed Model Name"
},
"DistributedValue": {
"input": {
"required": {
"default_value": [
"STRING",
{
"default": ""
}
],
"worker_values": [
"STRING",
{
"default": "{}"
}
]
},
"hidden": {
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"worker_id": [
"STRING",
{
"default": ""
}
]
}
},
"input_order": {
"required": [
"default_value",
"worker_values"
],
"hidden": [
"is_worker",
"worker_id"
]
},
"output": [
"*"
],
"output_name": [
"value"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "utils",
"display_name": "Distributed Value"
},
"ImageBatchDivider": {
"input": {
"required": {
"images": [
"IMAGE"
],
"divide_by": [
"INT",
{
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the batch into"
}
]
}
},
"input_order": {
"required": [
"images",
"divide_by"
]
},
"output": [
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*"
],
"output_name": [
"batch_1",
"batch_2",
"batch_3",
"batch_4",
"batch_5",
"batch_6",
"batch_7",
"batch_8",
"batch_9",
"batch_10"
],
"output_is_list": [
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
],
"is_input_list": false,
"output_node": true,
"category": "image",
"display_name": "Image Batch Divider"
},
"AudioBatchDivider": {
"input": {
"required": {
"audio": [
"AUDIO"
],
"divide_by": [
"INT",
{
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of sequential time segments to create"
}
]
}
},
"input_order": {
"required": [
"audio",
"divide_by"
]
},
"output": [
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*"
],
"output_name": [
"audio_1",
"audio_2",
"audio_3",
"audio_4",
"audio_5",
"audio_6",
"audio_7",
"audio_8",
"audio_9",
"audio_10"
],
"output_is_list": [
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
],
"is_input_list": false,
"output_node": true,
"category": "audio",
"display_name": "Audio Segment Divider"
},
"DistributedEmptyImage": {
"input": {
"required": {
"height": [
"INT",
{
"default": 64,
"min": 1,
"max": 4096,
"step": 1
}
],
"width": [
"INT",
{
"default": 64,
"min": 1,
"max": 4096,
"step": 1
}
],
"channels": [
"INT",
{
"default": 3,
"min": 1,
"max": 4,
"step": 1
}
]
}
},
"input_order": {
"required": [
"height",
"width",
"channels"
]
},
"output": [
"IMAGE"
],
"output_name": [
"IMAGE"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "image",
"display_name": "Distributed Empty Image"
},
"UltimateSDUpscaleDistributed": {
"input": {
"required": {
"upscaled_image": [
"IMAGE"
],
"model": [
"MODEL"
],
"positive": [
"CONDITIONING"
],
"negative": [
"CONDITIONING"
],
"vae": [
"VAE"
],
"seed": [
"INT",
{
"default": 0,
"min": 0,
"max": 18446744073709551615
}
],
"steps": [
"INT",
{
"default": 20,
"min": 1,
"max": 10000
}
],
"cfg": [
"FLOAT",
{
"default": 8.0,
"min": 0.0,
"max": 100.0
}
],
"sampler_name": [
[
"euler",
"euler_cfg_pp",
"euler_ancestral",
"euler_ancestral_cfg_pp",
"heun",
"heunpp2",
"exp_heun_2_x0",
"exp_heun_2_x0_sde",
"dpm_2",
"dpm_2_ancestral",
"lms",
"dpm_fast",
"dpm_adaptive",
"dpmpp_2s_ancestral",
"dpmpp_2s_ancestral_cfg_pp",
"dpmpp_sde",
"dpmpp_sde_gpu",
"dpmpp_2m",
"dpmpp_2m_cfg_pp",
"dpmpp_2m_sde",
"dpmpp_2m_sde_gpu",
"dpmpp_2m_sde_heun",
"dpmpp_2m_sde_heun_gpu",
"dpmpp_3m_sde",
"dpmpp_3m_sde_gpu",
"ddpm",
"lcm",
"ipndm",
"ipndm_v",
"deis",
"cfgpp_ud10_ab",
"res_multistep",
"res_multistep_cfg_pp",
"res_multistep_ancestral",
"res_multistep_ancestral_cfg_pp",
"gradient_estimation",
"gradient_estimation_cfg_pp",
"er_sde",
"seeds_2",
"seeds_3",
"sa_solver",
"sa_solver_pece",
"ddim",
"uni_pc",
"uni_pc_bh2"
]
],
"scheduler": [
[
"simple",
"sgm_uniform",
"karras",
"exponential",
"ddim_uniform",
"beta",
"normal",
"linear_quadratic",
"kl_optimal"
]
],
"denoise": [
"FLOAT",
{
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}
],
"tile_width": [
"INT",
{
"default": 512,
"min": 64,
"max": 2048,
"step": 8
}
],
"tile_height": [
"INT",
{
"default": 512,
"min": 64,
"max": 2048,
"step": 8
}
],
"padding": [
"INT",
{
"default": 32,
"min": 0,
"max": 256,
"step": 8
}
],
"mask_blur": [
"INT",
{
"default": 8,
"min": 0,
"max": 256
}
],
"force_uniform_tiles": [
"BOOLEAN",
{
"default": true
}
],
"tiled_decode": [
"BOOLEAN",
{
"default": false
}
]
},
"hidden": {
"multi_job_id": [
"STRING",
{
"default": ""
}
],
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"master_url": [
"STRING",
{
"default": ""
}
],
"enabled_worker_ids": [
"STRING",
{
"default": "[]"
}
],
"worker_id": [
"STRING",
{
"default": ""
}
],
"tile_indices": [
"STRING",
{
"default": ""
}
],
"dynamic_threshold": [
"INT",
{
"default": 8,
"min": 1,
"max": 64
}
]
}
},
"input_order": {
"required": [
"upscaled_image",
"model",
"positive",
"negative",
"vae",
"seed",
"steps",
"cfg",
"sampler_name",
"scheduler",
"denoise",
"tile_width",
"tile_height",
"padding",
"mask_blur",
"force_uniform_tiles",
"tiled_decode"
],
"hidden": [
"multi_job_id",
"is_worker",
"master_url",
"enabled_worker_ids",
"worker_id",
"tile_indices",
"dynamic_threshold"
]
},
"output": [
"IMAGE"
],
"output_name": [
"IMAGE"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "image/upscaling",
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
}
}
+607
View File
@@ -0,0 +1,607 @@
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
No managed ComfyUI service, user workflow, model, or remote worker is touched.
"""
import asyncio
from collections import Counter
import copy
import importlib
import json
import os
from pathlib import Path
import sys
import tempfile
import time
import uuid
from unittest.mock import patch
ROOT = Path(__file__).resolve().parents[2]
COMFY_ROOT = Path(sys.argv[1]).resolve()
WORKER = '--worker' in sys.argv
DELAY = float(sys.argv[3]) if WORKER else 0
sys.path.insert(0, str(COMFY_ROOT))
os.environ['COMFYUI_IS_WORKER'] = '1'
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
from app.assets.manager import default_asset_manager
import comfy.model_management
import execution
import nodes
import server
import torch
from aiohttp import web
import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
UTILITY_CALLS, ITEM_EXECUTIONS = [], []
ALLOW_MASTER = False
COUNT = 7
def source_rows():
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
for i in range(COUNT)]
class TestImageSource:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
OUTPUT_IS_LIST = (True,) * 6
FUNCTION = 'run'
def run(self):
assert not WORKER, 'upstream image source leaked into worker prompt'
SOURCE_CALLS.append(True)
rows = source_rows()
return tuple([row[key] for row in rows] for key in
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
class TestSharedFlag:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('BOOLEAN',)
FUNCTION = 'run'
def run(self):
assert not WORKER, 'shared value source leaked into worker prompt'
return (True,)
class TestCachedModel:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('MODEL',)
FUNCTION = 'run'
def run(self):
MODEL_CALLS.append(True)
return (object(),)
def assert_item_execution(prompt, index):
"""Prove the operation belongs to an item, not merely a participating host."""
from comfy_execution.utils import get_executing_context
context = get_executing_context()
assert context is not None
assert prompt['2']['inputs'].get('worker_item_token'), 'per-item node ran in outer graph'
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
assert len(running) == 1, running
entry = running[0]
if WORKER:
assert entry[1] == context.prompt_id and entry[2] == prompt
assert entry[4] == ['4'], entry[4]
else:
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
assert set(entry[4]) == {'4', '5'}, entry[4]
return {'index': index, 'prompt_id': context.prompt_id,
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
class TestRequiredOutputUtility:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
'hidden': {'prompt': 'PROMPT'}}
OUTPUT_NODE = True
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
def run(self, image, index, prompt):
assert_item_execution(prompt, index)
UTILITY_CALLS.append(index)
return (image,)
class TestDiscardedPreview:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, image):
raise AssertionError('discarded per-item output preview executed')
class TestImageOperation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
'text': ('STRING', {'forceInput': True}),
'strength': ('FLOAT', {'forceInput': True}),
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
'hidden': {'prompt': 'PROMPT'}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
BRANCH_CALLS.append(index)
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
until = time.monotonic() + DELAY
while time.monotonic() < until:
comfy.model_management.throw_exception_if_processing_interrupted()
time.sleep(0.01)
assert seed == 100 + index
assert flag is True
assert text == f'item-{index}'
assert strength == (index + 1) / 16
assert torch.all(reference_image == strength)
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
class TestCapture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
INPUT_IS_LIST = True
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
assert not WORKER, 'post-End capture leaked into a worker prompt'
CAPTURED.append((images, masks, texts))
return ()
async def setup():
assets = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), assets)
assert await nodes.load_custom_node(str(ROOT))
module = sys.modules[str(ROOT).replace('.', '_x_')]
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
TestRequiredOutputUtility=TestRequiredOutputUtility,
TestDiscardedPreview=TestDiscardedPreview)
# Match the CORS-enabled direct-worker launch contract for canvas probes.
app = web.Application(middlewares=[server.create_cors_middleware('*')])
# Register the actual feature routes, not test replacements.
app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
return ps, assets, module, app
async def listen(app):
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
return runner, site._server.sockets[0].getsockname()[1]
def finish_queue_item(ps, item_id, executor):
ps.prompt_queue.task_done(item_id, executor.history_result,
status=execution.PromptQueue.ExecutionStatus(
status_str='success' if executor.success else 'error',
completed=executor.success, messages=executor.status_messages),
process_item=lambda item: item[:5] + item[6:])
async def worker_main():
ps, assets, module, app = await setup()
wake = asyncio.Event()
interrupted_ids, submissions = [], []
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
async def prompt(request):
data = await request.json()
# Use core's actual validation/enqueue/response schema, including
# partial_execution_targets; only observe it and wake our CPU consumer.
response = await core_handlers['POST', '/prompt'](request)
if response.status == 200:
accepted = json.loads(response.body)
_, pending = ps.prompt_queue.get_current_queue_volatile()
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
'queued': entry[4], 'node_ids': sorted(entry[2]),
'end_ids': [key for key, node in entry[2].items()
if node['class_type'] == 'DistributedImageListEnd']})
wake.set()
return response
async def interrupt(request):
interrupted_ids.append((await request.json())['prompt_id'])
return await core_handlers['POST', '/interrupt'](request)
async def stats(request):
running, pending = ps.prompt_queue.get_current_queue_volatile()
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
async def consume():
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
while True:
await wake.wait()
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
entry, item_id = queued
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
wake.clear()
app.add_routes([route for route in ps.routes if
(route.path in ('/queue', '/history/{prompt_id}') or
(route.path == '/prompt' and route.method == 'GET'))])
app.router.add_post('/prompt', prompt)
app.router.add_post('/interrupt', interrupt)
app.router.add_get('/test-stats', stats)
runner, port = await listen(app)
consumer = asyncio.create_task(consume())
print(f'READY:{port}', flush=True)
try:
await asyncio.Event().wait()
finally:
consumer.cancel()
await asyncio.gather(consumer, return_exceptions=True)
await runner.cleanup()
def graph(callback, worker_ids):
return {
'1': {'class_type': 'TestImageSource', 'inputs': {}},
# Generic Autogrow value slots carry per-item reference images, text,
# strength and seed (one value per image), plus a shared boolean.
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'mask': ['1', 2],
'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
'3': {'class_type': 'TestImageOperation', 'inputs': {
key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
| {'image': ['8', 0], 'model': ['6', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
}
async def coordinator():
global ALLOW_MASTER
ps, assets, module, app = await setup()
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
config_mod = importlib.import_module(module.__name__ + '.utils.config')
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
runner, port = await listen(app)
ps.port = port # Match the listener used by local-worker callback resolution.
processes, stderr_files = [], []
# Old worker capability refusal is read-only and precedes input evaluation.
legacy_app = web.Application()
legacy_requests = []
async def legacy(request):
legacy_requests.append(request.method)
return web.json_response({'error': 'not found'}, status=404)
legacy_app.router.add_route('*', '/{path:.*}', legacy)
legacy_runner, legacy_port = await listen(legacy_app)
dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
try:
try:
await dispatch_module.preflight_image_list_workers([
{'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
raise AssertionError('legacy worker accepted')
except RuntimeError as exc:
assert 'upgrade' in str(exc), str(exc)
assert legacy_requests == ['GET'] and not SOURCE_CALLS
finally:
await legacy_runner.cleanup()
print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
scratch.mkdir(parents=True, exist_ok=True)
try:
with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
workers = []
for name, delay in [('fast', 0.02), ('slow', 0.5)]:
err = open(Path(tmp) / f'{name}.log', 'w+')
stderr_files.append(err)
proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
env={**os.environ, 'TMPDIR': tmp})
processes.append(proc)
async def ready():
while True:
line = await proc.stdout.readline()
if not line:
err.seek(0)
raise AssertionError(f'Worker exited: {err.read()}')
if line.startswith(b'READY:'):
return int(line.split(b':')[1])
worker_port = await asyncio.wait_for(ready(), 60)
workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
'port': worker_port, 'enabled': True})
config = {'workers': workers, 'settings': {'websocket_orchestration': False, 'master_delegate_only': True},
'master': {'host': f'http://127.0.0.1:{port}'}}
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
valid = await execution.validate_prompt('network-acceptance', prompt, None)
assert valid[0], valid
assert {'4', '5', '8', '9'} == set(valid[2]), valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
# Exercise production queue preparation; only master enqueue is redirected
# to this isolated executor, not the managed server's live queue.
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
prompt['2']['inputs'].pop('image_list_master_callback')
saved_prompt = copy.deepcopy(prompt)
outer_executions = []
async def enqueue(prepared, workflow, client_id, **kwargs):
targets = kwargs['partial_execution_targets']
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
assert planned == prepared and targets == planned_targets
if '9' in saved_prompt and 'TestImageOperation' in {
node['class_type'] for node in prepared.values()}:
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
assert set(targets) == {'4', '5'}, targets
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
checked = await execution.validate_prompt(pid, prepared, targets)
assert checked[0], checked
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
# must prove its roots from this real running record, not a mock.
number = len(outer_executions)
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
entry, item_id = ps.prompt_queue.get(timeout=0)
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
before_captures = len(CAPTURED)
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
finish_queue_item(ps, item_id, executor)
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
history = ps.prompt_queue.get_history(pid)[pid]
assert Counter(history['prompt'][4]) == Counter(targets)
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
outer_executions.append({'prompt_id': pid, 'targets': targets,
'captures': len(CAPTURED) - before_captures})
if executor.success:
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
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', [w['id'] for w in workers]), 45)
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
if not executor.success:
for err in stderr_files:
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS
assert len(CAPTURED) == 1
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
assert torch.equal(masks[i], row['mask'] * 0.5), i
assert texts[i] == f'item-{i}:seed-{100 + i}', texts
async with aiohttp.ClientSession() as session:
stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
stats.append(await response.json())
counts = [len(row['calls']) for row in stats]
assert counts[0] > counts[1] > 0, counts
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
for row in stats:
assert Counter(row['utility_calls']) == Counter(row['calls']), row
assert len(row['submissions']) == len(row['calls'])
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
for submission in row['submissions']:
assert submission['requested'] == submission['queued'] == ['4'], submission
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
assert '5' not in submission['node_ids']
assert not runtime._MASTER_JOBS
# Three sequential regions on one worker, including optional outputs
# left unconnected. Region N+1 must not race the prior worker End.
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
previous = '1'
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
'images': [previous, 0]}}
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': [start, 0], 'context': [start, 3]}}
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
saved_chained = copy.deepcopy(chained)
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
chained, None, 'test', ['slow']), 45)
assert chained == saved_chained
assert executor.success, executor.status_messages
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
assert CAPTURED[1][1:] == ([None, None], [None, None])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
# Same public queue path, now with master participation enabled. The
# parent queue stays occupied while its local slot processes items.
ALLOW_MASTER = True
config['settings']['master_delegate_only'] = False
participation_counts = {}
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)
async with aiohttp.ClientSession() as session:
before_worker_stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
before_worker_stats.append(await response.json())
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
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(
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'
local_calls = BRANCH_CALLS[before_local:]
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
assert local_calls, f'{mode}: master did not process any items'
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
remote_calls = []
async with aiohttp.ClientSession() as session:
for worker, before_stats in zip(workers, before_worker_stats):
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
current = await response.json()
calls = current['calls'][len(before_stats['calls']):]
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
submissions = current['submissions'][len(before_stats['submissions']):]
events = current['item_executions'][len(before_stats['item_executions']):]
assert Counter(utilities) == Counter(calls), (mode, current)
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
for submission in submissions:
assert submission['requested'] == submission['queued'] == ['4'], submission
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
assert '5' not in submission['node_ids']
remote_calls.extend(calls)
assert not current['running'] and not current['pending'], current
if mode.startswith('mixed'):
assert remote_calls, 'mixed execution failed to use direct workers'
else:
assert not remote_calls and len(local_calls) == COUNT
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
images, masks, texts = CAPTURED[-1]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
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) == 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'],
'required_utility_once_per_item': True, 'preview_discarded': True,
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
ALLOW_MASTER = False
config['settings']['master_delegate_only'] = True
# Preserve the original direct-runtime cancellation/failure fixture:
# it has no outer queue and must not claim target-control evidence.
direct_prompt = copy.deepcopy(prompt)
del direct_prompt['8'], direct_prompt['9']
direct_prompt['3']['inputs']['image'] = ['2', 0]
# Real Comfy user-interrupt flag while a real worker executor is active.
slow = workers[1]
async with aiohttp.ClientSession() as session:
async def slow_stats():
async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
return await response.json()
async def until(predicate):
for _ in range(150):
current = await slow_stats()
if predicate(current):
return current
await asyncio.sleep(0.02)
raise AssertionError('worker state transition timed out')
before = await until(lambda state: not state['running'] and not state['pending'])
context = runtime.create_master_job('2', '4', direct_prompt, source_rows(), ['slow'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
task = asyncio.create_task(runtime.run_image_list_job(context))
active = await until(lambda state: bool(state['running']))
owned_id = active['running'][0]
comfy.model_management.interrupt_current_processing(True)
try:
await asyncio.wait_for(task, 5)
raise AssertionError('interrupt was ignored')
except runtime.ImageListCancelledError:
pass
finally:
comfy.model_management.interrupt_current_processing(False)
after = await until(lambda state: not state['running'] and not state['pending'])
assert after['interrupted'] == [owned_id], after
assert len(after['calls']) <= len(before['calls']) + 1
assert set(before['history']) <= set(after['history'])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
# A worker exception before End must fail, not hang or return a partial list.
bad_rows = source_rows()
bad_rows[0]['seed'] = -1
context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
try:
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
raise AssertionError('worker exception was ignored')
except RuntimeError as exc:
assert 'failed image-list item' in str(exc), str(exc)
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
finally:
for proc in processes:
if proc.returncode is None:
proc.terminate()
await asyncio.gather(*(proc.wait() for proc in processes))
for err in stderr_files:
err.close()
await runner.cleanup()
network = importlib.import_module(module.__name__ + '.utils.network')
session = await network.get_client_session()
await session.close()
if __name__ == '__main__':
asyncio.run(worker_main() if WORKER else coordinator())
+469
View File
@@ -0,0 +1,469 @@
"""Real CPU Comfy V3 + HTTP proof; invoked only by the opt-in parent test.
The only harness pieces are deterministic source/assertion nodes and a bounded
queue worker. /prompt, /history and /view are Comfy's real handlers; compilation,
packing, staging, codecs, bridge execution, response decoding and Collector are
unmodified production code. No normal instance or Runpod client is contacted.
"""
import asyncio
import copy
import importlib
import io as bytes_io
import json
import os
from pathlib import Path
import socket
import sys
import types
import wave
from contextlib import asynccontextmanager
from fractions import Fraction
ROOT = Path(__file__).resolve().parents[2]
def bootstrap():
if len(sys.argv) != 3:
raise SystemExit("Use test_media_end_to_end.py with explicit COMFYUI_SOURCE_ROOT")
core, run_root = (Path(value).resolve() for value in sys.argv[1:])
assert run_root.is_dir() and run_root != core and not run_root.is_relative_to(core)
for name in ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
assert os.environ.get(name) == "", f"CPU isolation missing: {name}"
assert not any(name in sys.modules for name in ("torch", "execution", "folder_paths"))
sys.path.insert(0, str(core))
import comfy.cli_args
args = comfy.cli_args.args
# Set CPU and all writable roots BEFORE any framework/tensor import.
args.cpu = True
args.disable_assets = True
args.disable_all_custom_nodes = True
args.enable_manager = False
args.offline = True
args.base_directory = str(run_root)
args.front_end_root = str(run_root)
args.database_url = "sqlite:///:memory:"
for name in ("input", "output", "temp", "user"):
(run_root / name).mkdir()
setattr(args, name + "_directory", str(run_root / name))
return run_root
async def check(base):
import av
import numpy as np
import torch
import execution
import folder_paths
import nodes
import server
from aiohttp import web
from app.assets.manager import default_asset_manager
from comfy_api.v0_0_2 import io
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
from comfy_extras import nodes_audio
assert not torch.cuda.is_initialized()
torch.set_num_threads(1)
assert execution.comfy.model_management.get_torch_device().type == "cpu"
for name in ("input", "output", "temp", "user"):
assert Path(getattr(folder_paths, "get_" + name + "_directory")()).is_relative_to(base)
# Namespace scaffolding avoids the unrelated pack bootstrap (service/config
# side effects). Every imported module below is the real source on disk.
package_name = "_media_e2e_acceptance"
for suffix in ("", ".api", ".api.orchestration", ".nodes"):
package = types.ModuleType(package_name + suffix)
package.__path__ = [str(ROOT.joinpath(*suffix.lstrip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[package.__name__] = package
def load(name):
return importlib.import_module(package_name + "." + name)
asset_manager = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), asset_manager)
ps.distributed_pending_jobs = {}
ps.distributed_jobs_lock = asyncio.Lock()
compiler = load("api.orchestration.serverless_media_prompt")
packer = load("api.orchestration.serverless_typed_inputs")
decoder = load("api.orchestration.serverless_typed_results")
serializer = load("api.orchestration.runpod_client").serialize_request
handler = load("serverless_bridge.handler")
contract = load("serverless_bridge.contract")
video_codec = load("serverless_bridge.video_codec")
bridge = load("serverless_bridge")
collector = load("nodes.media_collector").DistributedMediaCollector
extension = await bridge.comfy_entrypoint()
bridge_classes = await extension.get_node_list()
assert {cls.GET_SCHEMA().node_id for cls in bridge_classes} == {"MediaImportV1", "MediaExportV1"}
sockets = list(contract.SOCKET_KINDS)
kinds = list(contract.SOCKET_KINDS.values())
assert kinds == ["IMAGE", "AUDIO", "MASK", "STRING", "VIDEO"]
io_types = [io.Image, io.Audio, io.Mask, io.String, io.Video]
fixtures, expected_sinks, seen_sinks, source_calls = {}, {}, {}, []
class AcceptanceMediaSource(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="AcceptanceMediaSource", inputs=[io.String.Input("key")],
outputs=[kind.Output(is_output_list=True) for kind in io_types])
@classmethod
def execute(cls, key):
source_calls.append(key)
return io.NodeOutput(*(fixtures[key][socket] for socket in sockets))
def tensor_equal(actual, expected):
assert actual.device.type == "cpu" and actual.dtype == expected.dtype
assert actual.shape == expected.shape
# Byte comparison preserves signed zero and catches dtype narrowing.
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def exact(kind, actual, expected):
if kind in ("IMAGE", "MASK"):
tensor_equal(actual, expected)
elif kind == "AUDIO":
assert actual["sample_rate"] == expected["sample_rate"]
tensor_equal(actual["waveform"], expected["waveform"])
elif kind == "STRING":
assert type(actual) is str and actual.encode("utf-8") == expected.encode("utf-8")
else:
assert video_codec.encode_video(actual) == expected
if type(actual) is VideoFromFile:
buffer = vars(actual)["_VideoFromFile__file"]
assert type(buffer) is bytes_io.BytesIO and not buffer.closed
position = buffer.tell()
buffer.seek(0)
with av.open(buffer) as container:
assert len(list(container.decode(video=0))) == 3
buffer.seek(position)
elif type(actual) is VideoFromList:
for item in actual.videos:
exact("VIDEO", item, video_codec.encode_video(item))
else:
assert type(actual) is VideoFromComponents
components = actual.get_components()
assert type(components.frame_rate) is Fraction
assert components.images.dtype == torch.float64
assert components.alpha.dtype == torch.bfloat16
assert components.audio["waveform"].dtype == torch.float64
assert actual.get_bit_depth() == 10 and actual.get_color_space() == "HDR PQ"
class AcceptanceMediaAssert(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(node_id="AcceptanceMediaAssert", inputs=[io.String.Input("tag"),
*[kind.Input(socket, optional=True, **({"force_input": True} if socket == "texts" else {}))
for socket, kind in zip(sockets, io_types)]], outputs=[],
is_input_list=True, is_output_node=True)
@classmethod
def execute(cls, tag, **media):
assert len(tag) == 1 and tag[0] not in seen_sinks
expected = expected_sinks[tag[0]]
assert set(media) == set(expected)
for socket in expected:
assert len(media[socket]) == len(expected[socket]), (tag, socket)
# Collector must preserve native objects and ordered item lists,
# not merge batches, resample, convert types, or encode again.
assert all(a is e for a, e in zip(media[socket], expected[socket])), (tag, socket)
counts = {socket: len(media[socket]) for socket in expected}
seen_sinks[tag[0]] = counts
return io.NodeOutput(ui={"accepted": [counts]})
for cls in [*bridge_classes, collector, AcceptanceMediaSource, AcceptanceMediaAssert]:
assert issubclass(cls, io.ComfyNode)
cls.GET_NODE_INFO_V1()
nodes.NODE_CLASS_MAPPINGS[cls.GET_SCHEMA().node_id] = cls
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def tiny_video(variant):
buffer = bytes_io.BytesIO()
with av.open(buffer, "w", format="matroska") as container:
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
stream.width = stream.height = 16
stream.pix_fmt = "yuv420p"
for level in (16 + variant, 64 + variant, 128 + variant):
frame = av.VideoFrame.from_ndarray(np.full((16, 16, 3), level, dtype=np.uint8), format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
return buffer.getvalue()
def native_fixtures(variant):
image = torch.arange(48, dtype=torch.float64).reshape(2, 2, 3, 4) / 997 + variant / 100
image[0, 0, 0, 0] = -0.0
mask = torch.tensor([[[-0.0, 0.125, 0.875]]], dtype=torch.bfloat16)
audio = {"waveform": torch.arange(38, dtype=torch.float64).reshape(1, 2, 19) / 997,
"sample_rate": 44100 + variant}
audio["waveform"][0, 0, 0] = -0.0
component = VideoFromComponents(VideoComponents(
images=image[..., :3], frame_rate=Fraction(30000, 1001), alpha=mask.reshape(1, 1, 3, 1).repeat(2, 2, 1, 1),
audio=audio, metadata={"title": "雪", "nested": [True, None, -4, 0.125]}),
bit_depth=10, color_space="HDR PQ")
payload = tiny_video(variant)
file = VideoFromFile(bytes_io.BytesIO(payload), start_time=-0.0, duration=0, crop=(0, 0, 8, 8))
sequence = VideoFromList([VideoFromFile(bytes_io.BytesIO(payload)), file], complete_audio=audio)
return {"images": [image, image[:1].clone()], "audio": [audio], "masks": [mask],
"texts": [f" café 🎥 {variant}\r\n雪 ", ""], "video": [component, sequence]}
for key, variant in (("local", 0), ("worker-a", 1), ("worker-b", 2)):
fixtures[key] = native_fixtures(variant)
native_root = base / "master-input"
native_root.mkdir()
wav = native_root / "stereo.wav"
pcm = np.array([-32768, -12001, -1, 0, 1, 8193, 32767, 17] * 2, dtype="<i2")
with wave.open(str(wav), "wb") as out:
out.setnchannels(2)
out.setsampwidth(2)
out.setframerate(22050)
out.writeframes(pcm.tobytes())
waveform, sample_rate = nodes_audio.load(str(wav))
clip = native_root / "native.mkv"
clip.write_bytes(tiny_video(3))
native_expected = copy.copy(fixtures["worker-a"])
native_expected["audio"] = [{"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}]
native_expected["video"] = [VideoFromFile(str(clip))]
expectations = {}
for key, values in (("worker-a", native_expected), ("worker-b", fixtures["worker-b"])):
expectations[key] = {socket: [video_codec.encode_video(v, allowed_roots=[native_root]) if socket == "video" else v
for v in values[socket]] for socket in sockets}
# Native routes only, bound to a new ephemeral loopback listener. No fake
# history, permissive validator, bridge/codec mock, or existing service.
allowed_paths = {"/prompt", "/history/{prompt_id}", "/view", "/queue"}
ps.app.add_routes([route for route in ps.routes if route.path in allowed_paths])
completed = []
stop = asyncio.Event()
async def queue_worker():
while not stop.is_set():
item = await asyncio.to_thread(ps.prompt_queue.get, 0.1)
if item is None:
continue
queued, item_id = item
assert len(queued) in (5, 6), "Unsupported core PromptQueue tuple; adapt acceptance worker explicitly"
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, queued[2], queued[1], queued[3], queued[4])
status = execution.PromptQueue.ExecutionStatus(
status_str="success" if executor.success else "error", completed=executor.success,
messages=executor.status_messages)
ps.prompt_queue.task_done(item_id, executor.history_result, status=status)
completed.append((queued, executor.success, executor.status_messages))
@asynccontextmanager
async def serve():
runner = web.AppRunner(ps.app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
ALLOWED_CONNECTIONS.add(("127.0.0.1", port))
worker = asyncio.create_task(queue_worker())
try:
yield f"http://127.0.0.1:{port}"
finally:
stop.set()
await asyncio.wait_for(worker, 10)
await runner.cleanup()
ALLOWED_CONNECTIONS.clear()
grouped_results = {}
async with serve() as url:
for key in ("worker-a", "worker-b"):
participant = "serverless:" + key
request_id = "acceptance-" + key
graph = {"source": node("AcceptanceMediaSource", key=key),
"collector": node("DistributedMediaCollector", load_balance=False,
multi_job_id="job-main", worker_id=participant,
**{socket: ["source", index] for index, socket in enumerate(sockets)}),
"text-collector": node("DistributedMediaCollector", load_balance=False,
multi_job_id="job-text", worker_id=participant, texts=["source", 3])}
if key == "worker-a":
graph.update({"wav": node("LoadAudio", audio=wav.name), "clip": node("LoadVideo", file=clip.name)})
graph["collector"]["inputs"].update(audio=["wav", 0], video=["clip", 0])
before = copy.deepcopy(graph)
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
assert graph == before
assert len(manifest) == 6 and {b["kind"] for b in manifest} == set(kinds)
folder_paths.set_input_directory(str(native_root))
workflow_before = copy.deepcopy(workflow)
images, artifacts, refs = await packer.prepare_media_inputs(workflow)
assert images == [] and workflow == workflow_before
if key == "worker-a":
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"} and len(artifacts) == 2
assert len(refs) == 2
assert {ref["input_name"] for ref in refs} == {"audio", "video"}
assert all(workflow[ref["node_id"]]["class_type"] == "MediaExportV1" for ref in refs)
# Endpoint cannot fall back to the master's local files/loaders.
wav.unlink()
clip.unlink()
print("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK WAV + native file VIDEO")
else:
assert not artifacts and not refs
folder_paths.set_input_directory(str(base / "input"))
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
"artifacts": artifacts, "inputs": refs, "outputs": manifest}
body = json.loads(serializer(workflow, images, 30, media=media))
assert decoder.validate_prepared_job({"worker": {"id": participant, "media_profile": contract.PROFILE,
"timeout_seconds": 30}, "workflow": workflow, "images": images,
"media": media, "manifest": manifest}) is True
before = copy.deepcopy(body["input"])
result = await handler.execute_media_request(body["input"], comfy_url=url,
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
assert body["input"] == before
assert completed and completed[-1][1], completed
queued = completed[-1][0]
assert set(queued[4]) == {b["node_id"] for b in manifest}
submitted = queued[2]
if key == "worker-a":
assert "wav" not in submitted and "clip" not in submitted
imports = [n for n in submitted.values() if n["class_type"] == "MediaImportV1"]
assert len(imports) == 2
assert {json.loads(n["inputs"]["artifact_json"])["kind"] for n in imports} == {"AUDIO", "VIDEO"}
assert not list((base / "input").rglob("*.bin"))
assert not list((base / "temp").rglob("*.bin"))
if key == "worker-b":
# Retire the originals as well as handler-owned staging/export
# files before constructing the result's lazy native videos.
for original in fixtures[key]["video"][1].videos:
vars(original)["_VideoFromFile__file"].close()
# Exercise the wire format, not an in-memory handler-only object.
result = json.loads(json.dumps(result, allow_nan=False))
grouped = await decoder.collect_typed_results(result, media, manifest)
assert set(grouped) == {"job-main", "job-text"}
assert set(grouped["job-main"]) == set(sockets)
assert grouped["job-text"]["texts"] == expectations[key]["texts"]
for socket, kind in zip(sockets, kinds):
values, expected = grouped["job-main"][socket], expectations[key][socket]
assert len(values) == len(expected), (key, socket)
for actual, wanted in zip(values, expected):
exact(kind, actual, wanted)
if key == "worker-b":
original = fixtures[key]["video"][0].get_components()
decoded = grouped["job-main"]["video"][0].get_components()
for field in ("images", "alpha"):
assert getattr(decoded, field).data_ptr() != getattr(original, field).data_ptr()
assert decoded.audio["waveform"].data_ptr() != original.audio["waveform"].data_ptr()
assert decoded.metadata == original.metadata and decoded.metadata is not original.metadata
grouped_results[participant] = grouped
# Late invalid artifact must reject the entire grouped response.
corrupt = copy.deepcopy(result)
corrupt["media"]["artifacts"][-1]["sha256"] = "0" * 64
try:
await decoder.collect_typed_results(corrupt, media, manifest)
except ValueError:
pass
else:
raise AssertionError("Result decoder accepted a corrupt final artifact")
assert ps.distributed_pending_jobs == {}
# Exercise user-selected image return encoding through the actual
# compiler, HTTP handler, V3 exporter and atomic response decoder.
for fmt in ("png", "jpeg", "h264"):
key = "encoded-" + fmt
fixtures[key] = native_fixtures(3)
# Keep full alpha precision and use nontrivial image dimensions.
torch.manual_seed(41)
rgba = torch.rand((2, 64, 64, 4), dtype=torch.float64)
fixtures[key]["images"] = [rgba]
request_id = "request-" + key
graph = {"source": node("AcceptanceMediaSource", key=key),
"collector": node("DistributedMediaCollector", transfer_format=fmt,
transfer_quality=73, multi_job_id="job-" + key,
worker_id="serverless:" + key,
**{socket: ["source", i] for i, socket in enumerate(sockets)})}
original = copy.deepcopy(graph)
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
"artifacts": [], "inputs": [], "outputs": manifest}
body = json.loads(serializer(workflow, [], 30, media=media))
result = await handler.execute_media_request(body["input"], comfy_url=url,
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
assert graph == original
grouped = await decoder.collect_typed_results(result, media, manifest)
actual = grouped["job-" + key]
image = actual["images"][0]
assert image.shape == rgba.shape and image.dtype == rgba.dtype
assert torch.equal(image[..., 3], rgba[..., 3]), "alpha must not be discarded/narrowed"
assert not torch.equal(image[..., :3], rgba[..., :3]), "selection must actually encode"
if fmt == "png":
pixels = (rgba[..., :3].float().clamp(0, 1) * 255).to(torch.uint8)
assert torch.equal(image[..., :3], (pixels.float() / 255).to(rgba.dtype))
for socket, kind in zip(sockets[1:], kinds[1:]):
for got, wanted in zip(actual[socket], fixtures[key][socket]):
exact(kind, got, video_codec.encode_video(wanted) if kind == "VIDEO" else wanted)
assert not list((base / "temp").rglob("*.bin"))
print("IMAGE_TRANSFER_HTTP_OK png jpeg h264; alpha and other media exact")
print("REAL_HTTP_BRIDGE_RESULTS_OK all five kinds; 6 bindings/request; native VIDEO remains readable after cleanup")
print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery")
participants = list(grouped_results)
for delegate in (False, True):
tag = "delegate" if delegate else "master"
job = "collect-" + tag
queue = asyncio.Queue()
for participant in reversed(participants):
await queue.put({"worker_id": participant, "media": grouped_results[participant]["job-main"], "is_last": True})
ps.distributed_pending_jobs[job] = queue
ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-main"] for p in participants]
expected_sinks[tag] = {socket: [item for values in ordered for item in values[socket]] for socket in sockets}
graph = {"collector": node("DistributedMediaCollector", load_balance=False, multi_job_id=job,
enabled_worker_ids=json.dumps(participants), expected_media=json.dumps(sockets), delegate_only=delegate),
"assert": node("AcceptanceMediaAssert", tag=tag,
**{socket: ["collector", index] for index, socket in enumerate(sockets)})}
if not delegate:
graph["source"] = node("AcceptanceMediaSource", key="local")
graph["collector"]["inputs"].update({socket: ["source", index] for index, socket in enumerate(sockets)})
# The second group from the very same complete response must reach its
# own Collector, rather than being folded into the first collector job.
text_job, text_tag = job + "-text", tag + "-text"
text_queue = asyncio.Queue()
for participant in reversed(participants):
await text_queue.put({"worker_id": participant,
"media": grouped_results[participant]["job-text"], "is_last": True})
ps.distributed_pending_jobs[text_job] = text_queue
text_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-text"] for p in participants]
expected_sinks[text_tag] = {"texts": [item for values in text_ordered for item in values["texts"]]}
graph["text-collector"] = node("DistributedMediaCollector", load_balance=False, multi_job_id=text_job,
enabled_worker_ids=json.dumps(participants), expected_media='["texts"]', delegate_only=delegate)
graph["text-assert"] = node("AcceptanceMediaAssert", tag=text_tag, texts=["text-collector", 3])
if not delegate:
graph["text-collector"]["inputs"]["texts"] = ["source", 3]
valid = await execution.validate_prompt(job, graph, None)
assert valid[0] and not valid[3], valid
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, job, {}, valid[2])
assert executor.success, executor.status_messages
assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]]
assert executor.history_result["outputs"]["text-assert"]["accepted"] == [seen_sinks[text_tag]]
assert job not in ps.distributed_pending_jobs and text_job not in ps.distributed_pending_jobs
assert source_calls == ["worker-a", "worker-b", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
assert set(seen_sinks) == {"master", "delegate", "master-text", "delegate-text"}
assert not list((base / "output").rglob("*"))
assert not torch.cuda.is_initialized()
print("MASTER_AND_DELEGATE_ORDER_OK", json.dumps(seen_sinks, sort_keys=True))
print("MEDIA_END_TO_END_CPU_OK real PromptExecutor/HTTP; no queue-orchestration or provider acceptance claimed")
ALLOWED_CONNECTIONS = set()
def network_guard(event, args):
"""Fail closed on provider/normal-instance traffic, including during imports."""
if event == "socket.connect":
sock, address = args
if sock.family in (socket.AF_INET, socket.AF_INET6):
assert address in ALLOWED_CONNECTIONS, "Acceptance forbids non-owned network connections"
elif event == "socket.bind":
sock, address = args
if sock.family in (socket.AF_INET, socket.AF_INET6):
assert address == ("127.0.0.1", 0), "Acceptance binds only ephemeral loopback"
if __name__ == "__main__":
sys.addaudithook(network_guard)
asyncio.run(asyncio.wait_for(check(bootstrap()), timeout=120))
+484
View File
@@ -0,0 +1,484 @@
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
No HTTP listener, workers, model downloads or live installation changes.
The original contracts were captured from 32ac027 using the same core.
"""
import asyncio
import inspect
import json
import os
from pathlib import Path
import sys
import tempfile
from unittest.mock import patch
from PIL import Image
ROOT = Path(__file__).resolve().parents[2]
COMFY_ROOT = Path(sys.argv[1]).resolve()
sys.path.insert(0, str(COMFY_ROOT))
os.environ['COMFYUI_IS_WORKER'] = '1'
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
from app.assets.manager import default_asset_manager
from comfy_api.v0_0_2 import io
import execution
import nodes
import server
import torch
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
# Inputs intentionally added after the V1 baseline was captured. They must be
# optional, appended after the baseline inputs, and keep legacy behaviour as
# their default so saved workflows are unaffected.
ADDED_OPTIONAL_INPUTS = {
'DistributedCollector': {
'transfer_format': [
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
{'default': 'legacy_png'},
],
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
},
}
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd', 'DistributedMediaCollector'}
def normalized_input(value):
kind = value[0]
opts = dict(value[1]) if len(value) > 1 else {}
if kind == 'STRING':
# V3 serializes the same default single-line widget explicitly.
opts.setdefault('multiline', False)
if kind == 'COMBO':
kind = opts['options']
opts = {key: val for key, val in opts.items() if key != 'options'}
if isinstance(kind, list):
# Single-selection is the V1 dropdown default as well.
opts.setdefault('multiselect', False)
return [kind, opts]
async def map_node(cls, values, extra=None):
"""Use the actual input parser and sanitized V3 executor clones."""
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
assert not missing, missing
# Use a worker thread for synchronous nodes: collector/upscale bridge back
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
def run():
return asyncio.run(execution._async_map_node_over_list(
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
return await asyncio.to_thread(run)
def check_saved_workflows(mapping):
checked = 0
occurrences = 0
for path in sorted((ROOT / 'workflows').glob('*.json')):
workflow = json.loads(path.read_text())
local_nodes = {node['id']: node for node in workflow['nodes']}
for saved in workflow['nodes']:
if saved['type'] not in mapping:
continue
occurrences += 1
info = mapping[saved['type']].GET_NODE_INFO_V1()
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
for socket in saved.get('inputs', []):
assert socket['name'] in inputs, (path.name, saved['id'], socket)
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
for index, output in enumerate(saved.get('outputs', [])):
assert output['type'] == info['output'][index], (path.name, index)
for link in workflow['links']:
if link[1] == saved['id']:
assert link[2] < len(info['output']), (path.name, link)
target = local_nodes[link[3]]['inputs'][link[4]]
assert target['type'] == info['output'][link[2]], (path.name, link)
checked += 1
assert checked == 5 and occurrences == 10, (checked, occurrences)
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
async def check_execution(mapping, module, prompt_server, asset_manager):
result = await map_node(mapping['DistributedSeed'],
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
assert result[0].result == (126,)
values = json.dumps({'_type': 'INT', '3': '17'})
result = await map_node(mapping['DistributedValue'],
{'default_value': '4', 'worker_values': values,
'is_worker': True, 'worker_id': 'worker_2'})
assert result[0].result == (17,)
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
{'extra_pnginfo': metadata})
assert result[0].result == ('model.ckpt',)
assert result[0].ui == {'text': ['model.ckpt']}
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
result = await map_node(mapping['DistributedEmptyImage'],
{'height': 8, 'width': 8, 'channels': 3})
empty = result[0].result[0]
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
for count in (1, 3, 10):
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
assert len(result[0].result) == 10
assert torch.equal(torch.cat(result[0].result[:count]), images)
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
assert len(result[0].result) == 10
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
collector = mapping['DistributedCollector']
prepared, missing, hidden = execution.get_input_data(
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
assert not missing
prepared['images'] = [images[:2], images[2:]]
prepared['audio'] = [audio, audio]
result = await execution._async_map_node_over_list(
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
assert torch.equal(result[0].result[0], images)
assert result[0].result[1]['waveform'].shape[-1] == 66
# Exercise real async aggregation without an HTTP endpoint or worker.
queue = asyncio.Queue()
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
result = await map_node(collector,
{'images': images[:2], 'load_balance': False,
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
assert torch.equal(result[0].result[0], images[:3])
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
# The registered node must forward the transfer options to the runtime
# helper, and default to the legacy transfer when they are not set.
collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
forwarded = []
def fake_run(self, **kwargs):
forwarded.append(kwargs)
return (kwargs['images'], kwargs['audio'])
with patch.object(collector_runtime, 'run', fake_run):
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
'transfer_format': 'av1', 'transfer_quality': 70})
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
# Set values arrive list-wrapped (is_input_list); unset ones use the
# execute() defaults. The runtime helper unwraps both forms.
assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
(['av1'], [70]), ('legacy_png', 90)], forwarded
# Prove V3's per-execution class clones do not leak helper instance state;
# replace only the GPU/model boundary, not the input parser or executor.
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
seen = []
def fake_upscale(self, *args):
assert not hasattr(self, 'acceptance_marker')
self.acceptance_marker = True
seen.append(args)
return (args[0],)
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
'force_uniform_tiles': True, 'tiled_decode': False,
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
with patch.object(upscale_runtime, 'run', fake_upscale):
for _ in range(2):
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
assert result[0].result[0] is images
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
'tile_indices', 'dynamic_threshold')), seen[0]
import math
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
'pass_through': True}},
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
}
valid = await execution.validate_prompt('v3-graph', graph, None)
assert valid[0], valid
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
assert not rejected[0]
errors = [error for entry in rejected[3].values() for error in entry['errors']]
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
import folder_paths
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
assert executor.success, executor.status_messages
history = executor.history_result
record = history['outputs']['4']['images'][0]
preview = Path(temp) / record.get('subfolder', '') / record['filename']
with Image.open(preview) as image:
assert image.size == (8, 8) and image.mode == 'RGB'
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
# A worker prompt terminated by the sink validates and runs through the
# real executor without writing anything to the temp directory.
sink_graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 4, 'color': 0}},
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
'pass_through': True}},
'3': {'class_type': 'DistributedWorkerSink', 'inputs': {'value': ['2', 0]}},
}
valid = await execution.validate_prompt('v3-sink', sink_graph, None)
assert valid[0], valid
assert valid[2] == ['3'], valid[2]
collected = []
def record_run(self, **kwargs):
collected.append(kwargs['images'])
return (kwargs['images'][0], None)
with tempfile.TemporaryDirectory(prefix='v3-sink-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp), patch.object(collector_runtime, 'run', record_run):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, sink_graph, 'v3-sink', {}, valid[2])
assert executor.success, executor.status_messages
assert len(collected) == 1 and collected[0][0].shape == (4, 8, 8, 3), collected
assert list(Path(temp).rglob('*')) == [], list(Path(temp).rglob('*'))
assert executor.history_result['outputs'].get('3', {}) == {}, executor.history_result
print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing')
async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
"""Execute real native V3 list/lazy flow; master must skip the processing branch."""
branch_calls = []
class ImageListProbe:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
CATEGORY = 'tests'
def run(self, image):
branch_calls.append(image)
return (image,)
nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
runtime = sys.modules[module.__name__ + '.image_list.runtime']
original_contexts = set(runtime._MASTER_JOBS)
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
async def fake_worker_results(_context_id, *, model_callback=None):
state = runtime.get_master_job(_context_id)
assert state.prompt_id == 'image-list-lazy', state.prompt_id
assert state.workflow_meta == workflow_meta, state.workflow_meta
assert len(prompt_server.app.cleanup_ctx) >= 1, 'lifecycle cleanup was not registered'
return (
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
[None, None, None], [None, None, None],
)
graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'image_list_enabled_worker_ids': '["test-worker"]',
'image_list_master_callback': 'http://127.0.0.1:1',
'image_list_workflow_meta': workflow_meta,
}},
'3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'context': ['2', 3],
}},
'5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
}
valid = await execution.validate_prompt('image-list-lazy', graph, None)
assert valid[0], valid
validator = runtime.validate_image_list_regions
assert 'texts' not in graph['4']['inputs']
assert len(validator(graph)) == 1
connected_text_graph = json.loads(json.dumps(graph))
connected_text_graph['6'] = {
'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 4]},
}
connected_text_graph['4']['inputs']['texts'] = ['6', 0]
assert len(validator(connected_text_graph)) == 1
for literal_text in ('', 'per-item text'):
literal_text_graph = json.loads(json.dumps(graph))
literal_text_graph['4']['inputs']['texts'] = literal_text
try:
validator(literal_text_graph)
except ValueError:
pass
else:
raise AssertionError(f'End accepted unconnected literal text: {literal_text!r}')
print('IMAGE_LIST_END_TEXT_INPUT_OK native forceInput; unconnected omitted; linked accepted; literals rejected')
import folder_paths
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
try:
# Resolve genuine Comfy output/temp annotations against the pinned core,
# not only the controlled folder_paths seam in unit tests.
media = sys.modules[module.__name__ + '.api.orchestration.media_sync']
with tempfile.TemporaryDirectory(prefix='media-roots-', dir=scratch) as media_root:
for kind in ('input', 'output', 'temp'):
source = Path(media_root) / kind / 'shared/ref.png'
source.parent.mkdir(parents=True)
source.write_bytes(kind.encode())
with patch.object(folder_paths, 'input_directory', str(Path(media_root) / 'input')), \
patch.object(folder_paths, 'output_directory', str(Path(media_root) / 'output')), \
patch.object(folder_paths, 'temp_directory', str(Path(media_root) / 'temp')):
for kind in ('output', 'temp'):
refs = media._find_media_references({'1': {'inputs': {'image': f'shared/ref.png [{kind}]'}}})
assert len(refs) == 1
assert media._load_media_file_sync(refs[0])[0] == kind.encode()
with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp), \
patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
assert executor.success, executor.status_messages
assert branch_calls == [], 'master evaluated the lazy per-image branch'
records = executor.history_result['outputs']['5']['images']
assert len(records) == 3, records
for record in records:
preview = Path(temp) / record.get('subfolder', '') / record['filename']
with Image.open(preview) as image:
assert image.size == (8, 8) and image.mode == 'RGB'
finally:
nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
for context_id in set(runtime._MASTER_JOBS) - original_contexts:
runtime.remove_master_job(context_id)
print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews verified')
async def main():
asset_manager = default_asset_manager()
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
module = sys.modules[str(ROOT).replace('.', '_x_')]
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
extension = await module.comfy_entrypoint()
classes = await extension.get_node_list()
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES)
assert set(mapping) == set(BASELINE) | ADDED_NODES
for node_id in ADDED_NODES:
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
media_info = mapping['DistributedMediaCollector'].GET_NODE_INFO_V1()
assert media_info['output'] == ['IMAGE', 'AUDIO', 'MASK', 'STRING', 'VIDEO'], media_info
assert all(media_info['output_is_list']), media_info
assert media_info['input']['optional']['texts'][1]['forceInput'] is True
sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
assert sink_info['output_node'] and sink_info['dev_only'], sink_info
for node_id, cls in mapping.items():
if node_id in ADDED_NODES:
continue
assert issubclass(cls, io.ComfyNode)
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
old = dict(BASELINE[node_id])
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
# V1's ByPassTypeTuple advertises '*' when indexed, while its
# underlying tuple and existing frontend declare IMAGE/AUDIO.
# Native V3 declares all ten existing typed sockets explicitly.
assert old['output'] == ['*'] * 10
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
for group in ('required', 'optional'):
old_inputs = old['input'].get(group, {})
new_inputs = dict(new['input'].get(group, {}))
if group == 'optional':
for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
assert name in new_inputs, (node_id, name, 'added input missing')
actual = normalized_input(new_inputs.pop(name))
kind, opts = normalized_input(expected)
assert actual[0] == kind, (node_id, name, actual)
for key, value in opts.items():
assert actual[1].get(key) == value, (node_id, name, key, actual)
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
for name, original in old_inputs.items():
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
assert old[key] == new[key], (node_id, key, old[key], new[key])
# Standard context lives in cls.hidden; orchestrator metadata remains
# accepted by its original kwarg name, without creating new widgets.
signature = inspect.signature(cls.execute)
for name, field in old['input'].get('hidden', {}).items():
if isinstance(field, list):
assert cls.GET_SCHEMA().accept_all_inputs, node_id
assert name in signature.parameters, (node_id, name)
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
else:
assert name in new['input']['hidden'], (node_id, name)
expected_hidden = []
if node_id == 'DistributedModelName':
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
if old['output_node']:
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
assert start_schema.node_id == 'DistributedImageListStart'
assert start_schema.display_name == 'Distributed Image List Start'
assert end_schema.node_id == 'DistributedImageListEnd'
assert end_schema.display_name == 'Distributed Image List End'
texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
texts_schema = texts_input.as_dict()
assert texts_schema.get('forceInput') is True, texts_schema
assert 'default' not in texts_schema, texts_schema
assert [out.io_type for out in start_schema.outputs] == [
'IMAGE', 'MASK', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + ['*'] * 10
assert [out.display_name for out in start_schema.outputs][4:] == [f'value_{i}' for i in range(1, 11)]
assert all(out.is_output_list for out in start_schema.outputs)
assert start_schema.is_input_list and start_schema.accept_all_inputs
start_info = json.loads(json.dumps(mapping['DistributedImageListStart'].GET_NODE_INFO_V1()))
assert list(start_info['input']['optional']) == ['images', 'mask', 'values'], start_info['input']
values_spec = start_info['input']['optional']['values']
assert values_spec[0] == 'COMFY_AUTOGROW_V3', values_spec
template = values_spec[1]['template']
assert template['names'] == [f'value_{i}' for i in range(1, 11)] and template['min'] == 0, template
assert template['input']['required']['value'][0] == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN', template
# Autogrow + INPUT_IS_LIST through the real parser/list mapper, values
# still split 1-or-N, unconnected slots ignored.
# get_input_data must accept Autogrow names; constants are wrapped once.
parsed, missing, hidden = execution.get_input_data(
{'values.value_1': 'shared', 'values.value_3': 4},
mapping['DistributedImageListStart'], unique_id='autogrow-probe', extra_data={})
assert not missing and parsed['values.value_1'] == ['shared'], parsed
# Linked inputs arrive as the upstream output list (one element per value).
prepared = {'images': [torch.zeros((3, 2, 2, 3))], 'values.value_1': ['shared'],
'values.value_3': [4, 5, 6]}
seen = {}
runtime_mod = sys.modules[module.__name__ + '.image_list.runtime']
def capture(_images, _mask, values):
seen['values'] = values
raise RuntimeError('captured')
nodes_mod = sys.modules[module.__name__ + '.nodes.image_list']
with patch.object(nodes_mod, 'normalize_image_list_inputs', capture):
try:
await asyncio.to_thread(lambda: asyncio.run(execution._async_map_node_over_list(
'autogrow-probe', 'autogrow-probe', mapping['DistributedImageListStart'], prepared,
'execute', v3_data=hidden)))
except RuntimeError as exc:
assert 'captured' in str(exc), exc
assert seen['values'] == {'value_1': ['shared'], 'value_3': [4, 5, 6]}, seen
rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
assert all(out.is_output_list for out in end_schema.outputs)
assert end_schema.is_input_list and end_schema.is_output_node
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
print('SCHEMA_PARITY_OK', len(mapping))
check_saved_workflows(mapping)
await check_execution(mapping, module, prompt_server, asset_manager)
await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
print('V3_ACCEPTANCE_OK')
asyncio.run(main())
+146
View File
@@ -0,0 +1,146 @@
import importlib.util
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
class _PromptQueue:
def __init__(self):
self.items = []
def put(self, item):
self.items.append(item)
def _load_async_helpers_module():
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
package_name = "dist_async_helpers_testpkg"
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
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{package_name}.utils"] = utils_pkg
execution_module = types.ModuleType("execution")
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
return (True, None, ["9"], {})
execution_module.validate_prompt = _validate_prompt
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
sys.modules["execution"] = execution_module
prompt_server = types.SimpleNamespace(
trigger_on_prompt=lambda payload: payload,
number=12,
prompt_queue=_PromptQueue(),
)
server_module = types.ModuleType("server")
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
sys.modules["server"] = server_module
network_module = types.ModuleType(f"{package_name}.utils.network")
network_module.get_server_loop = lambda: None
sys.modules[f"{package_name}.utils.network"] = network_module
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", 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)
return module, prompt_server
async_helpers, prompt_server = _load_async_helpers_module()
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
result = await async_helpers.queue_prompt_payload(
{"1": {"class_type": "Node"}},
workflow_meta={"id": "workflow-1"},
client_id="client-1",
include_queue_metadata=True,
)
self.assertIsInstance(result["prompt_id"], str)
self.assertTrue(result["prompt_id"])
self.assertEqual(result["number"], 12)
self.assertEqual(result["node_errors"], {})
self.assertEqual(prompt_server.number, 13)
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
queued_item = prompt_server.prompt_queue.items[0]
self.assertEqual(queued_item[0], 12)
extra_data = queued_item[3]
self.assertEqual(extra_data["client_id"], "client-1")
self.assertIn("create_time", extra_data)
self.assertIsInstance(extra_data["create_time"], int)
self.assertGreater(extra_data["create_time"], 0)
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.module, self.server = _load_async_helpers_module()
self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
self.module.execution.validate_prompt = self.validate
async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
seen = []
def callback(payload):
seen.append(payload["partial_execution_targets"][:])
payload["partial_execution_targets"].append("10")
return {"prompt": payload["prompt"]}
self.server.trigger_on_prompt = callback
targets = ["9"]
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
self.assertEqual(seen, [["9"]])
self.assertEqual(targets, ["9"])
self.assertEqual(self.validate.call_args.args[2], ["9"])
self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])
async def test_invalid_requests_never_reach_validation_or_queue(self):
for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
with self.subTest(targets=targets):
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
self.validate.assert_not_awaited()
self.assertEqual(self.server.prompt_queue.items, [])
async def test_removed_target_after_callback_is_rejected(self):
self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
self.validate.assert_not_awaited()
self.assertEqual(self.server.prompt_queue.items, [])
async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
(["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
(["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
with self.subTest(roots=roots, errors=errors):
self.validate.return_value = (True, None, roots, errors)
with self.assertRaises(self.module.PromptValidationError):
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
self.assertEqual(self.server.prompt_queue.items, [])
async def test_none_preserves_legacy_payload_and_validation(self):
seen = []
self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
await self.module.queue_prompt_payload(self.prompt)
self.assertNotIn("partial_execution_targets", seen[0])
self.assertIsNone(self.validate.call_args.args[2])
if __name__ == "__main__":
unittest.main()
+638
View File
@@ -0,0 +1,638 @@
import importlib.util
import sys
import types
import asyncio
from pathlib import Path
import pytest
import torch
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[1] / "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_collector_module():
# Restore shared dependencies even when a different test imports this loader.
# Keep real modules imported during loading (notably extension modules) cached.
with pytest.MonkeyPatch.context() as monkeypatch:
return _load_collector_module_with_stubs(monkeypatch)
def _load_collector_module_with_stubs(monkeypatch):
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
package_name = "dist_collector_list_testpkg"
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
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
nodes_pkg.__path__ = []
sys.modules[f"{package_name}.nodes"] = nodes_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)
class _Routes:
def post(self, _path):
return lambda fn: fn
def get(self, _path):
return lambda fn: fn
prompt_server = types.SimpleNamespace(
routes=_Routes(),
distributed_jobs_lock=None,
distributed_pending_jobs={},
)
server_module = types.ModuleType("server")
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
monkeypatch.setitem(sys.modules, "server", server_module)
comfy_module = types.ModuleType("comfy")
model_management = types.ModuleType("comfy.model_management")
class InterruptProcessingException(Exception):
pass
model_management.InterruptProcessingException = InterruptProcessingException
model_management.throw_exception_if_processing_interrupted = lambda: None
comfy_module.model_management = model_management
comfy_utils = types.ModuleType("comfy.utils")
class ProgressBar:
def __init__(self, _total):
self.total = _total
self.updates = []
def update(self, value):
self.updates.append(value)
comfy_utils.ProgressBar = ProgressBar
comfy_module.utils = comfy_utils
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
aiohttp_module = types.ModuleType("aiohttp")
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_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
config_module = types.ModuleType(f"{package_name}.utils.config")
config_module.get_worker_timeout_seconds = lambda: 0.1
config_module.load_config = lambda: {"workers": []}
config_module.is_master_delegate_only = lambda: False
sys.modules[f"{package_name}.utils.config"] = config_module
constants_module = types.ModuleType(f"{package_name}.utils.constants")
constants_module.HEARTBEAT_INTERVAL = 1.0
constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
sys.modules[f"{package_name}.utils.constants"] = constants_module
image_module = types.ModuleType(f"{package_name}.utils.image")
def _ensure_contiguous(tensor):
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
image_module.ensure_contiguous = _ensure_contiguous
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
image_module.pil_to_tensor = lambda value: value
sys.modules[f"{package_name}.utils.image"] = image_module
network_module = types.ModuleType(f"{package_name}.utils.network")
network_module.build_worker_url = lambda worker: "http://worker"
network_module.get_client_session = lambda: None
network_module.probe_worker = lambda *_args, **_kwargs: None
sys.modules[f"{package_name}.utils.network"] = network_module
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
audio_payload_module.encode_audio_payload = lambda audio: audio
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", 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)
return module
def test_collector_opts_into_comfyui_list_inputs():
collector = _load_collector_module().DistributedCollectorNode
assert collector.INPUT_IS_LIST is True
def test_collector_exposes_images_as_optional_input():
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
assert "images" not in input_types["required"]
assert input_types["optional"]["images"] == ("IMAGE",)
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
collector = _load_collector_module().DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
images, returned_audio = collector.run(images=None, audio=[audio])
assert images is None
assert returned_audio is audio
def test_collector_rejects_missing_images_and_audio():
collector = _load_collector_module().DistributedCollectorNode()
try:
collector.run(images=None, audio=None)
except ValueError as exc:
assert "image or audio" in str(exc).lower()
else:
raise AssertionError("Expected collector to reject a run with no media input")
def test_delegate_only_master_allows_no_local_media_input():
collector = _load_collector_module().DistributedCollectorNode()
images, audio = collector.run(
images=None,
audio=None,
multi_job_id=["delegate-audio-job"],
delegate_only=[True],
enabled_worker_ids=["[]"],
)
assert images is None
assert tuple(audio["waveform"].shape) == (1, 2, 1)
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
collector = _load_collector_module().DistributedCollectorNode()
first = torch.zeros(1, 2, 2, 3)
second = torch.ones(1, 2, 2, 3)
images, audio = collector.run(
images=[first, second],
load_balance=[False],
audio=[None],
multi_job_id=[""],
is_worker=[False],
master_url=[""],
enabled_worker_ids=["[]"],
worker_batch_size=[1],
worker_id=[""],
pass_through=[False],
delegate_only=[False],
)
assert tuple(images.shape) == (2, 2, 2, 3)
assert torch.equal(images[0:1], first)
assert torch.equal(images[1:2], second)
assert tuple(audio["waveform"].shape) == (1, 2, 1)
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
first = torch.zeros(1, 2, 2, 3)
second = torch.ones(1, 2, 2, 3)
posted_payloads = []
class _FakeImage:
def save(self, fp, format=None, compress_level=None):
fp.write(b"png-bytes")
class _FakeResponse:
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
def raise_for_status(self):
return None
class _FakeSession:
def post(self, url, json, timeout):
posted_payloads.append(json)
return _FakeResponse()
async def _fake_get_client_session():
return _FakeSession()
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
module.get_client_session = _fake_get_client_session
module.encode_audio_payload = lambda _audio: None
images, audio = collector.run(
images=[first, second],
load_balance=[False],
audio=[None],
multi_job_id=["job-list-1"],
is_worker=[True],
master_url=["http://master"],
enabled_worker_ids=["[]"],
worker_batch_size=[1],
worker_id=["worker-a"],
pass_through=[False],
delegate_only=[False],
)
assert tuple(images.shape) == (2, 2, 2, 3)
assert tuple(audio["waveform"].shape) == (1, 2, 1)
assert len(posted_payloads) == 2
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
def test_audio_only_worker_sends_one_completion_without_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
posted = []
class _FakeResponse:
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
def raise_for_status(self):
return None
class _FakeSession:
def post(self, url, json, timeout):
posted.append((url, json, timeout.total))
return _FakeResponse()
async def _fake_get_client_session():
return _FakeSession()
module.get_client_session = _fake_get_client_session
module.encode_audio_payload = lambda value: {"encoded": value is audio}
images, returned_audio = collector.run(
images=None,
audio=[audio],
multi_job_id=["audio-job"],
is_worker=[True],
master_url=["http://master"],
worker_id=["worker-a"],
)
assert images is None
assert returned_audio is audio
assert len(posted) == 1
assert posted[0][0] == "http://master/distributed/job_complete"
assert posted[0][2] == 600
assert posted[0][1] == {
"job_id": "audio-job",
"worker_id": "worker-a",
"batch_idx": 0,
"audio": {"encoded": True},
"is_last": True,
}
def test_audio_only_master_combines_local_and_worker_audio():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=master_audio,
multi_job_id="audio-job",
enabled_worker_ids='["worker-a"]',
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=None,
multi_job_id="delegate-audio-job",
enabled_worker_ids='["worker-a"]',
delegate_only=True,
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
# --------------------------------------------------------------------------- #
# Compressed worker -> master transfer
# --------------------------------------------------------------------------- #
class _RecordingResponse:
def __init__(self, status=200, text=""):
self.status = status
self._text = text
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def text(self):
return self._text
def raise_for_status(self):
if self.status >= 400:
raise RuntimeError(f"HTTP {self.status}")
class _RecordingSession:
def __init__(self, statuses=None, text="Invalid JSON payload"):
self.calls = []
self._statuses = list(statuses or [])
self._text = text
def post(self, url, timeout, json=None, data=None):
self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
status = self._statuses.pop(0) if self._statuses else 200
return _RecordingResponse(status=status, text=self._text)
def _install_fake_transport(module, session):
async def _fake_get_client_session():
return session
module.get_client_session = _fake_get_client_session
# Record the multipart content instead of building a real aiohttp.FormData.
module.DistributedCollectorNode._build_multipart = staticmethod(
lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
)
def _run_worker(collector, images, audio=None, **kwargs):
return collector.run(
images=[images],
audio=[audio],
multi_job_id=["job-c"],
is_worker=[True],
master_url=["http://master"],
worker_id=["worker-a"],
**{key: [value] for key, value in kwargs.items()},
)
def test_collector_exposes_transfer_widgets_with_legacy_default():
module = _load_collector_module()
optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
formats, options = optional["transfer_format"]
assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
assert options["default"] == "legacy_png"
assert optional["transfer_quality"][0] == "INT"
def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
images = torch.rand(6, 32, 48, 3)
_run_worker(
collector,
images,
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="ffv1",
transfer_quality=90,
)
assert len(session.calls) == 1
call = session.calls[0]
assert call["url"] == "http://master/distributed/job_complete"
assert call["json"] is None
envelope = call["data"]["envelope"]
assert envelope["format"] == "ffv1"
assert (envelope["start_index"], envelope["count"]) == (0, 6)
assert (envelope["width"], envelope["height"]) == (48, 32)
assert envelope["is_last"] is True
assert envelope["audio"] == {"encoded": True}
assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
assert len(call["data"]["parts"]) == 1
def test_worker_marks_only_final_chunk_as_last():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.TRANSFER_MAX_CHUNK_BYTES = 1 # force one frame per request
_run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
envelopes = [call["data"]["envelope"] for call in session.calls]
assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
(0, 1, False),
(1, 1, False),
(2, 1, True),
]
assert "audio" not in envelopes[-1]
def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
from PIL import Image
import numpy as np
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession(statuses=[400])
_install_fake_transport(module, session)
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
)
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
assert session.calls[0]["data"] is not None # compressed attempt
legacy = [call["json"] for call in session.calls[1:]]
assert [payload["batch_idx"] for payload in legacy] == [0, 1]
assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
assert [payload["is_last"] for payload in legacy] == [False, True]
def test_worker_uses_legacy_path_for_unknown_transfer_format():
from PIL import Image
import numpy as np
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
)
_run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
assert len(session.calls) == 1
assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
def test_worker_falls_back_when_video_encoder_missing():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
codec = module.transfer_codec
original = codec.video_codec_available
codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
try:
_run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
finally:
codec.video_codec_available = original
assert session.calls[0]["data"]["envelope"]["format"] == "h264"
def test_worker_does_not_fall_back_on_non_legacy_rejection():
import pytest
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
_install_fake_transport(module, session)
with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
# One compressed attempt, no legacy resend of the batch.
assert len(session.calls) == 1
assert session.calls[0]["json"] is None
def test_worker_sends_nothing_when_a_frame_cannot_fit():
import numpy as np
import pytest
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
# Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
images = torch.zeros(3, 128, 128, 3)
images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
with pytest.raises(module.transfer_codec.FrameTooLargeError):
_run_worker(collector, images, transfer_format="png")
assert session.calls == []
def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
_run_worker(
collector,
torch.zeros(2, 16, 16, 3),
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="png",
)
multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
assert [e["is_last"] for e in multipart] == [False]
assert all("audio" not in e for e in multipart)
final = session.calls[-1]["json"]
assert final["is_last"] is True and final["audio"]["data"].startswith("x")
assert "image" not in final
def test_worker_attaches_audio_to_last_chunk_when_it_fits():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
session = _RecordingSession()
_install_fake_transport(module, session)
module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
_run_worker(
collector,
torch.zeros(2, 16, 16, 3),
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
transfer_format="png",
)
assert len(session.calls) == 1
envelope = session.calls[0]["data"]["envelope"]
assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}
+1 -1
View File
@@ -18,7 +18,7 @@ def _load_config_module():
del sys.modules[mod_name]
root_pkg = types.ModuleType(package_name)
root_pkg.__path__ = []
root_pkg.__path__ = [str(module_path.parent)]
sys.modules[package_name] = root_pkg
logging_module = types.ModuleType(f"{package_name}.logging")
+322 -2
View File
@@ -4,7 +4,7 @@ import sys
import types
import unittest
from pathlib import Path
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
def _load_dispatch_module():
@@ -55,7 +55,13 @@ def _load_dispatch_module():
def __init__(self, total=None):
self.total = total
class _ClientConnectorError(Exception):
class _ClientError(Exception):
pass
class _ClientConnectorError(_ClientError):
pass
class _ClientConnectionError(_ClientError):
pass
class _WSMsgType:
@@ -76,7 +82,9 @@ def _load_dispatch_module():
self.closed = True
aiohttp_module.ClientTimeout = _ClientTimeout
setattr(aiohttp_module, "ClientError", _ClientError)
aiohttp_module.ClientConnectorError = _ClientConnectorError
setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
aiohttp_module.WSMsgType = _WSMsgType
aiohttp_module.TCPConnector = _TCPConnector
aiohttp_module.ClientSession = _ClientSession
@@ -100,6 +108,113 @@ def _load_dispatch_module():
dispatch = _load_dispatch_module()
def _load_prompt_transform_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
package_name = "prompt_transform_worker_selection_testpkg"
for suffix in ("", ".api", ".api.orchestration", ".utils"):
package = types.ModuleType(package_name + suffix)
package.__path__ = []
sys.modules[package_name + suffix] = package
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
sys.modules[f"{package_name}.utils.logging"] = logging_module
spec = importlib.util.spec_from_file_location(
f"{package_name}.api.orchestration.prompt_transform",
module_path,
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
prompt_transform = _load_prompt_transform_module()
class _FakeResponse:
def __init__(self, status=200, payload=None, error=None):
self.status = status
self.payload = payload
self.error = error
async def __aenter__(self):
if self.error is not None:
raise self.error
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.payload
class _FakeSession:
def __init__(self, responses):
self.responses = responses
def get(self, url, timeout=None):
del timeout
worker_id = url.removeprefix("http://").split("/", 1)[0]
endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
response = self.responses[(worker_id, endpoint)]
if isinstance(response, BaseException):
return _FakeResponse(error=response)
status, payload = response
return _FakeResponse(status=status, payload=payload)
def _valid_capability():
return {
"supported": True,
"execution_targets": True,
"protocol_version": 2,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
}
def _valid_queue(running=None, pending=None):
return {"queue_running": running or [], "queue_pending": pending or []}
class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
async def test_http_preserves_targets_and_omits_default_field(self):
response = types.SimpleNamespace(raise_for_status=lambda: None)
class Response:
async def __aenter__(self): return response
async def __aexit__(self, *args): pass
calls = []
def post(url, **kwargs):
calls.append(kwargs['json'])
return Response()
session = types.SimpleNamespace(post=post)
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
partial_execution_targets=['9'])
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None)
self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
self.assertNotIn('partial_execution_targets', calls[1])
async def test_websocket_sends_targets_all_the_way_to_wire(self):
import json
sent = []
async def receive():
return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({
'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'],
'ok': True, 'prompt_id': 'worker-prompt'}))
ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append),
receive=receive, close=AsyncMock())
session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws))
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
use_websocket=True, return_prompt_id=True, partial_execution_targets=['9'])
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True)
self.assertEqual(result, 'worker-prompt')
self.assertEqual(sent[0]['partial_execution_targets'], ['9'])
self.assertNotIn('partial_execution_targets', sent[1])
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
async def test_select_active_workers_filters_offline(self):
workers = [
@@ -249,6 +364,211 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
self.assertIsNone(selected)
async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "offline"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "offline")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "offline")
})
responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "slow"}]
responses: dict = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "slow")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "slow")
})
responses[("slow", "queue")] = asyncio.TimeoutError()
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_excludes_busy_workers(self):
workers = [{"id": "idle"}, {"id": "busy"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("idle", "busy")
}
responses.update({
("idle", "queue"): (200, _valid_queue()),
("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
})
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"idle": True})
async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
workers = [{"id": "offline-a"}, {"id": "offline-b"}]
responses = {
(worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
for worker_id in ("offline-a", "offline-b")
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
self.assertEqual(
await dispatch.preflight_image_list_workers(workers, require_idle=False),
{},
)
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await dispatch.preflight_image_list_workers(workers, require_idle=True)
async def test_image_list_preflight_requires_exact_execution_targets_support(self):
for value in (None, False, 1, "true", []):
with self.subTest(value=value):
capability = _valid_capability()
if value is None:
capability.pop("execution_targets")
else:
capability["execution_targets"] = value
session = _FakeSession({("worker", "capabilities"): (200, capability)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "execution_targets.*upgrade"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
incompatible = _valid_capability()
incompatible["protocol_version"] = 1
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
incompatible = _valid_capability()
incompatible["node_ids"] = ["DistributedImageListStart"]
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
responses = {
("worker", "capabilities"): (200, _valid_capability()),
("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(ValueError, "programming error"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
class ImageListDirectWorkerResolutionTests(unittest.TestCase):
def setUp(self):
self.config = {"workers": [
{"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
{"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
{"id": "serverless", "enabled": True, "type": "serverless"},
{"id": "other", "enabled": True, "host": "other-host", "port": 8190},
]}
def test_implicit_selection_uses_only_enabled_direct_workers(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
def test_explicit_selection_includes_exact_disabled_direct_worker(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
self.assertEqual([worker["id"] for worker in workers], ["disabled"])
def test_explicit_empty_selection_returns_no_workers(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
def test_unknown_explicit_id_raises_instead_of_falling_back(self):
with self.assertRaisesRegex(RuntimeError, "missing"):
prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
def test_explicit_serverless_selection_is_excluded(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
if __name__ == "__main__":
unittest.main()
+341
View File
@@ -0,0 +1,341 @@
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
import asyncio
import unittest
import torch
from image_list.runtime import (
END_CLASS,
IMAGE_OUT,
CONTEXT_OUT,
START_CLASS,
VALUE_OUT_BASE,
ImageListCancelledError,
build_worker_prompt,
create_master_job,
decode_item,
decode_result,
dynamic_schedule,
encode_item,
encode_result,
normalize_image_list_inputs,
normalize_worker_result,
remove_master_job,
set_master_result,
validate_image_list_regions,
)
def start_node(image_link):
return {"class_type": START_CLASS, "inputs": {"images": image_link}}
def graph_one_region():
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
"3": start_node(["2", 0]),
"4": {"class_type": "GenericFrameOperation", "inputs": {
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", VALUE_OUT_BASE],
}},
"5": {"class_type": END_CLASS, "inputs": {
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
}},
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
}
def graph_two_sequential_regions():
return {
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
"2": start_node(["1", 0]),
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
"5": start_node(["4", 0]),
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
}
class ImageListNormalizationTests(unittest.TestCase):
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
images = [
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
]
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
rows = normalize_image_list_inputs(images, masks, {
"value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"],
"value_3": [0.2], "value_4": [37],
})
self.assertEqual(len(rows), 2)
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
self.assertTrue(torch.equal(rows[0]["values"]["value_1"], refs[:1]))
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
self.assertEqual([row["index"] for row in rows], [0, 1])
self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"])
self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2])
self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37])
def test_one_value_is_shared_and_n_values_pair_in_order(self):
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
rows = normalize_image_list_inputs(image, values={
"value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]})
self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3)
self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3)
self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3)
self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3)
rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]})
self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}])
with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"):
normalize_image_list_inputs(image, values={"value_1": ["one", "two"]})
with self.assertRaisesRegex(ValueError, "mask has 2 values"):
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)))
shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4)))
self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask))
def test_value_batches_stay_whole_and_only_images_split(self):
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
pair = torch.rand((2, 4, 4, 3)) # e.g. two references for one item
rows = normalize_image_list_inputs(images, values={"value_1": [pair]})
self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows))
per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))]
rows = normalize_image_list_inputs(images, values={"value_1": per_item})
self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1]))
def test_rejects_unsupported_values_and_counts(self):
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
with self.assertRaisesRegex(ValueError, "value_1 has 3 values"):
normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]})
with self.assertRaisesRegex(TypeError, "unsupported type dict"):
normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]})
with self.assertRaisesRegex(TypeError, "floating-point"):
normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]})
with self.assertRaisesRegex(TypeError, "finite"):
normalize_image_list_inputs(images, values={"value_1": [float("nan")]})
with self.assertRaisesRegex(ValueError, "Unknown Start value slot"):
normalize_image_list_inputs(images, values={"value_11": [1]})
with self.assertRaisesRegex(ValueError, "at least one image"):
normalize_image_list_inputs(None)
# Unconnected Autogrow slots arrive as None and are ignored.
rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None})
self.assertEqual([row["values"] for row in rows], [{}, {}])
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
"values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16),
"value_2": "safe text", "value_3": 0.625, "value_4": 42,
"value_5": False},
"index": 3}
restored = decode_item(encode_item(item))
self.assertTrue(torch.equal(restored["image"], item["image"]))
self.assertEqual(restored["image"].dtype, torch.float64)
self.assertTrue(torch.equal(restored["values"]["value_1"], item["values"]["value_1"]))
self.assertEqual(restored["values"]["value_1"].dtype, torch.float16)
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
self.assertEqual(restored["values"]["value_2"], "safe text")
self.assertEqual(restored["values"]["value_3"], 0.625)
self.assertIs(restored["values"]["value_4"], 42)
self.assertIs(restored["values"]["value_5"], False)
with self.assertRaisesRegex(ValueError, "unknown slot"):
decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}})
with self.assertRaisesRegex(ValueError, "unsupported type"):
decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}})
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
from image_list.runtime import encode_value
encode_value("/untrusted/path", "PATH")
def test_result_transport_preserves_ragged_result_shapes(self):
result = normalize_worker_result(
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
[torch.ones((1, 3, 5), dtype=torch.float16)],
["done"],
)
encoded = encode_result(result)
decoded = decode_result(encoded)
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
self.assertEqual(decoded["image"].dtype, torch.float32)
self.assertEqual(decoded["mask"].dtype, torch.float16)
self.assertEqual(decoded["text"], "done")
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
class ImageListGraphBoundaryTests(unittest.TestCase):
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
prompt = graph_one_region()
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
"ctx", 1, "http://master.invalid")
self.assertEqual(set(worker), {"1", "3", "4", "5"})
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
self.assertNotIn("images", worker["3"]["inputs"])
def test_worker_prompt_drops_per_item_value_links(self):
prompt = graph_one_region()
prompt["8"] = {"class_type": "StringList", "inputs": {}}
prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]})
import sys, types
from unittest.mock import patch
registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})}
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012", "ctx", 0, "http://m")
self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")])
self.assertNotIn("8", worker)
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
self.assertNotIn("6", worker)
def test_rejects_upstream_processing_and_leaf_image_bypasses(self):
# Both a leaf source and a generated batch must cross Start exactly once.
for generated in (False, True):
with self.subTest(generated=generated):
prompt = graph_one_region()
if generated:
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
prompt["4"]["inputs"]["reference"] = ["2", 0]
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
import sys
import types
from unittest.mock import patch
prompt = graph_one_region()
prompt["7"] = {"class_type": "SeedConstant", "inputs": {"value": 9}}
prompt["2"]["inputs"].update({"model": ["1", 0], "seed": ["7", 0]})
prompt["4"]["inputs"]["shared_seed"] = ["7", 0]
registry = {
"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL", "CLIP", "VAE")}),
"SeedConstant": type("Constant", (), {"RETURN_TYPES": ("INT",)}),
}
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token", "ctx", 0, "http://master")
self.assertNotIn("2", worker)
self.assertTrue({"1", "7"}.issubset(worker))
def test_unknown_or_list_producing_shared_dependency_is_not_exempt(self):
import sys
import types
from unittest.mock import patch
for node_class in (None, type("ListLoader", (), {
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,)})):
prompt = graph_one_region()
prompt["2"]["inputs"]["model"] = ["1", 0]
registry = {} if node_class is None else {"CheckpointLoaderSimple": node_class}
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
with self.assertRaisesRegex(ValueError, "bypasses Start.*1"):
validate_image_list_regions(prompt)
def test_two_sequential_regions_are_valid_and_isolated(self):
regions = validate_image_list_regions(graph_two_sequential_regions())
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
p = graph_two_sequential_regions()
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
"http://master.invalid")
self.assertEqual(set(second), {"5", "6", "7"})
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
prompt = graph_one_region()
prompt["5"]["inputs"]["context"] = ["2", 0]
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "typed context output"):
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
nested = graph_one_region()
nested["7"] = start_node(["4", 0])
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
nested["5"]["inputs"]["images"] = ["8", 0]
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
validate_image_list_regions(nested, explicit_targets=explicit_targets)
outside = graph_one_region()
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
for explicit_targets in (False, True):
with self.subTest(explicit_targets=explicit_targets):
with self.assertRaisesRegex(ValueError, "cross-boundary"):
validate_image_list_regions(outside, explicit_targets=explicit_targets)
def test_rejects_duplicate_or_missing_results(self):
prompt = graph_one_region()
prompt["5"]["inputs"].pop("images")
with self.assertRaisesRegex(ValueError, "images must be connected"):
validate_image_list_regions(prompt)
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
async def test_fast_workers_get_more_items_and_order_is_stable(self):
assignments = {"fast": [], "slow": []}
async def dispatch(worker, index):
assignments[worker].append(index)
delay = 0.002 if worker == "fast" else 0.025
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
return {"index": index, "task": task}
async def wait_next(inflight):
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
task = next(iter(done))
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
index, result = task.result()
return worker, index, result
async def cleanup(_worker, handle):
handle["task"].cancel()
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
async def dispatch(worker, index):
return {"index": index, "prompt_id": f"{worker}-{index}"}
cancelled = []
async def cleanup(worker, handle):
cancelled.append((worker, handle["prompt_id"]))
async def timeout(_inflight):
raise TimeoutError("worker timed out")
with self.assertRaisesRegex(TimeoutError, "timed out"):
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
cancelled.clear()
async def wrong_index(inflight):
worker, handle = next(iter(inflight.items()))
return worker, handle["index"] + 99, "bad"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
self.assertEqual(len(cancelled), 1)
cancelled.clear()
with self.assertRaises(ImageListCancelledError):
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
interrupted=lambda: True)
self.assertEqual(len(cancelled), 0)
async def test_no_workers_and_duplicate_results_fail_closed(self):
async def unused(*_args):
self.fail("must not dispatch without workers")
with self.assertRaisesRegex(RuntimeError, "No free"):
await dynamic_schedule(2, [], unused, unused, unused)
async def dispatch(worker, index):
return {"index": index}
async def duplicate(inflight):
worker, _handle = next(iter(inflight.items()))
return worker, 0, "same"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
if __name__ == "__main__":
unittest.main()
+245
View File
@@ -0,0 +1,245 @@
"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution."""
import copy
import sys
import types
import pytest
from image_list.execution_plan import prepare_image_list_execution
from image_list.runtime import (
CONTEXT_OUT,
END_CLASS,
IMAGE_OUT,
START_CLASS,
build_worker_prompt,
validate_image_list_regions,
)
class OutputNode:
OUTPUT_NODE = True
class SchemaOutputNode:
OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative.
@classmethod
def GET_NODE_INFO_V1(cls):
return {"output_node": True, "output": ("IMAGE",)}
class ModelLoader:
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
@pytest.fixture
def registry(monkeypatch):
classes = {
END_CLASS: SchemaOutputNode,
"SaveImage": OutputNode,
"PreviewImage": OutputNode,
"OutputUtility": OutputNode,
"SchemaUtility": SchemaOutputNode,
"ModelLoader": ModelLoader,
}
monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes))
return classes
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def graph():
return {
"loader": node("ModelLoader"),
"source": node("ImageSource"),
"start": node(START_CLASS, images=["source", 0]),
"process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]),
"end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]),
"save": node("SaveImage", images=["end", 0]),
}
def test_discards_multistep_previews_but_preserves_other_output_roots(registry):
prompt = graph()
prompt.update({
"debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]),
"debug-preview": node("PreviewImage", images=["debug-transform", 0]),
"process-preview": node("PreviewImage", images=["process", 0]),
"unrelated-source": node("ImageSource"),
"unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]),
"upstream-preview": node("PreviewImage", images=["source", 0]),
"post-transform": node("FrameOperation", images=["end", 0]),
"post-preview": node("PreviewImage", images=["post-transform", 0]),
})
before = copy.deepcopy(prompt)
with pytest.raises(ValueError, match="cross-boundary"):
validate_image_list_regions(prompt)
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"}
assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"]
validate_image_list_regions(prepared, explicit_targets=True)
assert prompt == before
@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"])
def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type):
prompt = graph()
prompt.update({
"utility-config": node("Constant", value=7),
"utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]),
"utility-second": node(class_type, images=["utility-first", 0]),
"result": node("FrameOperation", image=["utility-second", 0]),
"discarded-preview": node("PreviewImage", images=["utility-first", 0]),
})
prompt["end"]["inputs"]["images"] = ["result", 0]
before = copy.deepcopy(prompt)
registered = dict(registry)
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {"discarded-preview"}
assert targets == ["end", "save"]
with pytest.raises(ValueError, match="Output node"):
validate_image_list_regions(prepared)
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid")
assert set(worker) == set(prepared) - {"source", "save"}
assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]}
assert prompt == before
assert registry == registered
assert OutputNode.OUTPUT_NODE is True
assert SchemaOutputNode.OUTPUT_NODE is False
assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True
@pytest.mark.parametrize("source", ["start", "process"])
@pytest.mark.parametrize("direct", [False, True])
def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct):
prompt = graph()
prompt["post-process"] = node("Composite", images=["end", 0])
if not direct:
prompt["bypass"] = node("FrameOperation", image=[source, 0])
source = "bypass"
prompt["post-process"]["inputs"]["reference"] = [source, 0]
prompt["save"]["inputs"]["images"] = ["post-process", 0]
before = copy.deepcopy(prompt)
with pytest.raises(ValueError, match="cross-boundary"):
prepare_image_list_execution(prompt)
assert prompt == before
def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry):
prompt = graph()
prompt.update({
"next-start": node(START_CLASS, images=["end", 0]),
"next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]),
"next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]),
"next-save": node("SaveImage", images=["next-end", 0]),
"first-debug": node("PreviewImage", images=["process", 0]),
"second-debug": node("PreviewImage", images=["next-process", 0]),
})
renamed = {key: f"group:7/{key}" for key in prompt}
for spec in prompt.values():
for name, value in spec["inputs"].items():
if isinstance(value, list):
spec["inputs"][name] = [renamed[value[0]], value[1]]
prompt = {renamed[key]: spec for key, spec in prompt.items()}
prepared, targets = prepare_image_list_execution(prompt)
assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]}
assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")]
regions = validate_image_list_regions(prepared, explicit_targets=True)
assert {(r["start_id"], r["end_id"]) for r in regions} == {
(renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]),
}
second = next(r for r in regions if r["start_id"] == renamed["next-start"])
worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid")
assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")}
def test_shared_loader_and_its_independent_output_root_are_preserved(registry):
registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {})
prompt = graph()
prompt["source"]["inputs"]["model"] = ["loader", 0]
prompt["independent-save"] = node("SaveImage", model=["loader", 0])
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == prompt
assert targets == ["loader", "end", "save", "independent-save"]
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
assert "loader" in region["node_ids"]
assert "loader" not in region["branch_ids"]
@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), {
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,),
})])
def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class):
if unsafe_class is None:
registry.pop("ModelLoader")
else:
registry["ModelLoader"] = unsafe_class
prompt = graph()
prompt["source"]["inputs"]["model"] = ["loader", 0]
prompt["debug"] = node("PreviewImage", images=["process", 0])
with pytest.raises(ValueError, match="bypasses Start.*loader"):
prepare_image_list_execution(prompt)
def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry):
prompt = graph()
prompt["process"]["inputs"]["reference"] = ["source", 0]
prompt["debug"] = node("PreviewImage", images=["process", 0])
with pytest.raises(ValueError, match="bypasses Start.*source"):
prepare_image_list_execution(prompt)
@pytest.mark.parametrize("nested", [False, True])
def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested):
prompt = graph()
prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0])
prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0])
prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT])
with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"):
prepare_image_list_execution(prompt)
@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"])
def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid):
prompt = graph()
if invalid == "missing-end":
del prompt["end"]
elif invalid == "wrong-context":
prompt["end"]["inputs"]["context"] = ["source", 0]
elif invalid == "missing-images":
del prompt["end"]["inputs"]["images"]
else:
prompt["end"]["inputs"]["images"] = ["source", 0]
before = copy.deepcopy(prompt)
with pytest.raises(ValueError):
prepare_image_list_execution(prompt)
assert prompt == before
def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry):
prompt = graph()
prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]}
before = copy.deepcopy(prompt)
prepared, targets = prepare_image_list_execution(prompt)
prepared_again, targets_again = prepare_image_list_execution(prepared)
assert prepared_again == prepared
assert targets_again == targets
prepared["start"]["inputs"]["images"][0] = "changed"
prepared["process"]["_meta"]["tags"].append("execution-only")
assert prompt == before
assert prepared_again == before
def test_v3_false_output_flag_overrides_legacy_true(registry):
registry["NotOutput"] = type("NotOutput", (), {
"OUTPUT_NODE": True,
"GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}),
})
prompt = graph()
prompt["not-output"] = node("NotOutput", images=["end", 0])
prepared, targets = prepare_image_list_execution(prompt)
assert prepared == prompt
assert targets == ["end", "save"]
+189
View File
@@ -0,0 +1,189 @@
"""Regression coverage at the scheduler and real HTTP transport boundaries."""
import asyncio
import importlib.util
import json
from pathlib import Path
import sys
import types
import unittest
from unittest.mock import patch
import torch
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from image_list import runtime
from test_image_list import graph_two_sequential_regions, start_node
def load_routes():
root = Path(__file__).resolve().parents[1]
name = '_image_list_transport_tests'
routes = web.RouteTableDef()
modules = {
name: types.ModuleType(name),
name + '.api': types.ModuleType(name + '.api'),
name + '.image_list': types.ModuleType(name + '.image_list'),
name + '.image_list.runtime': runtime,
'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
instance=types.SimpleNamespace(routes=routes))),
}
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)
with patch.dict(sys.modules, modules):
spec.loader.exec_module(module)
return module, routes
class CapabilityTests(unittest.IsolatedAsyncioTestCase):
async def test_capability_advertises_target_control_without_tensor_protocol_bump(self):
module, _routes = load_routes()
response = await module.image_list_capabilities_endpoint(None)
payload = json.loads(response.text)
self.assertIs(payload['execution_targets'], True)
self.assertEqual(payload['protocol_version'], 2)
class TransportTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.module, routes = load_routes()
app = web.Application()
app.add_routes(routes)
self.client = TestClient(TestServer(app))
await self.client.start_server()
async def asyncTearDown(self):
await self.client.close()
runtime._WORKER_ITEMS.clear()
async def test_fragmented_chunked_item_upload_is_consumable(self):
token = runtime.create_worker_task_token()
item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
async def chunks():
yield data[:8]
await asyncio.sleep(0.02)
yield data[8:]
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 200, await response.text())
result = runtime.consume_worker_item(token)
self.assertTrue(torch.equal(result['image'], item['image']))
self.assertEqual(result['index'], 2)
async def test_complete_item_upload_and_replay_rejection(self):
token = runtime.create_worker_task_token()
data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
{'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 200, await response.text())
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 400)
runtime.consume_worker_item(token)
with self.assertRaisesRegex(RuntimeError, 'already consumed'):
runtime.consume_worker_item(token)
async def test_chunked_body_limit_is_aggregate(self):
async def chunks():
yield b'{"context_id":"'
await asyncio.sleep(0.01)
yield b'x' * 60
yield b'"}'
with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 400)
self.assertIn('exceeds', await response.text())
async def test_invalid_result_shapes_are_rejected(self):
for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
'mask': torch.zeros((1, 2, 3))})
result[kind]['shape'] = list(shape)
with self.assertRaises(ValueError):
runtime.decode_result(result)
class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
entered = asyncio.Event()
cancelled = asyncio.Event()
wait_cancelled = asyncio.Event()
dispatches, cleanups = [], []
async def dispatch(worker, index):
dispatches.append(index)
return {'index': index}
async def wait_next(_inflight):
entered.set()
try:
await asyncio.Event().wait()
finally:
wait_cancelled.set()
async def cleanup(worker, handle):
cleanups.append((worker, handle['index']))
task = asyncio.create_task(runtime.dynamic_schedule(
3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
await entered.wait()
cancelled.set()
with self.assertRaises(runtime.ImageListCancelledError):
await asyncio.wait_for(task, timeout=1)
self.assertTrue(wait_cancelled.is_set())
self.assertEqual(dispatches, [0])
self.assertEqual(cleanups, [('one', 0)])
async def test_already_interrupted_does_not_submit_any_work(self):
async def unused(*args):
self.fail('pre-cancelled job must not dispatch or wait')
with self.assertRaises(runtime.ImageListCancelledError):
await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
async def test_task_cancel_during_submit_waits_for_owned_handle(self):
entered, release = asyncio.Event(), asyncio.Event()
cleanups = []
async def dispatch(worker, index):
entered.set()
await release.wait()
return {'index': index, 'prompt_id': 'owned'}
async def unused(*args):
self.fail('cancelled dispatch must not enter result wait')
async def cleanup(worker, handle):
cleanups.append(handle['prompt_id'])
task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
await entered.wait()
task.cancel()
release.set()
with self.assertRaises(asyncio.CancelledError):
await task
self.assertEqual(cleanups, ['owned'])
def test_output_nodes_inside_region_are_rejected(self):
prompt = graph_two_sequential_regions()
output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
with patch.dict(sys.modules, {'nodes': registry}):
with self.assertRaisesRegex(ValueError, 'must be placed after'):
runtime.validate_image_list_regions(prompt)
def test_shared_scalars_are_validated_before_dispatch(self):
for value in (float('nan'), object(), torch.zeros(2, dtype=torch.int32)):
with self.assertRaises(TypeError):
runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), values={'value_1': [value]})
def test_three_sequential_regions_can_share_worker_loader(self):
prompt = graph_two_sequential_regions()
prompt.pop('8')
prompt.update({
'8': start_node(['7', 0]),
'9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
'10': {'class_type': runtime.END_CLASS,
'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
'11': {'class_type': 'GenericModelLoader', 'inputs': {}},
})
for node_id in ['3', '6', '9']:
prompt[node_id]['inputs']['model'] = ['11', 0]
# Shared dependencies are verified against the registered output schema.
loader = type('GenericModelLoader', (), {'RETURN_TYPES': ('MODEL',)})
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'GenericModelLoader': loader})
with patch.dict(sys.modules, {'nodes': registry}):
regions = runtime.validate_image_list_regions(prompt)
self.assertEqual(len(regions), 3)
for region in regions:
self.assertIn('11', region['node_ids'])
+254
View File
@@ -0,0 +1,254 @@
"""Tests for periodic image-list server lifecycle cleanup."""
import asyncio
import time
import unittest
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import patch
from aiohttp import web
from image_list import lifecycle, runtime
class ExecutionTargetProofTests(unittest.TestCase):
def setUp(self):
from test_image_list_master_runtime import graph
self.prompt = graph()
self.registry = SimpleNamespace(NODE_CLASS_MAPPINGS={
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})})
self.modules = patch.dict('sys.modules', {'nodes': self.registry})
self.modules.start()
self.addCleanup(self.modules.stop)
def server(self, roots, *, prompt=None, running=True):
item = (0, 'owned', self.prompt if prompt is None else prompt, {}, roots, {})
return SimpleNamespace(prompt_queue=SimpleNamespace(
get_current_queue_volatile=lambda: ([item], []) if running else ([], [item])))
def test_outer_roots_in_running_queue_are_authoritative(self):
self.assertTrue(lifecycle.has_safe_execution_targets(self.server(['end']), 'owned', self.prompt))
for roots in (['op', 'end'], [], ['end', 'end'], None, [True]):
with self.subTest(roots=roots):
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(roots), 'owned', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end']), 'other', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], running=False), 'owned', self.prompt))
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], prompt={}), 'owned', self.prompt))
def test_prompt_marker_is_not_an_execution_guarantee(self):
self.prompt['start']['inputs']['image_list_explicit_targets'] = True
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['op', 'end']), 'owned', self.prompt))
with self.assertRaisesRegex(ValueError, 'must be placed after'):
runtime.create_master_job('start', 'end', self.prompt, [], [], '', master_participates=True)
self.assertEqual(runtime._MASTER_JOBS, {})
def test_start_node_uses_running_queue_proof(self):
import importlib.util
import sys
from pathlib import Path
from types import ModuleType
import torch
root = Path(__file__).resolve().parents[1]
package = '_image_list_start_targets'
modules = {}
for suffix in ('', '.nodes', '.image_list', '.utils'):
module = ModuleType(package + suffix)
module.__path__ = [str(root / suffix.lstrip('.'))]
modules[package + suffix] = module
modules[package + '.image_list.runtime'] = runtime
modules[package + '.image_list.lifecycle'] = lifecycle
modules[package + '.utils.async_helpers'] = SimpleNamespace(run_async_in_server_loop=None)
io = SimpleNamespace(ComfyNode=type('ComfyNode', (), {}), NodeOutput=lambda *args: args,
**{name: object() for name in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean')})
modules['comfy_api.v0_0_2'] = SimpleNamespace(io=io)
modules['comfy_execution.utils'] = SimpleNamespace(
get_executing_context=lambda: SimpleNamespace(prompt_id='owned'))
server = SimpleNamespace(PromptServer=SimpleNamespace(instance=self.server(['end'])))
modules['server'] = server
spec = importlib.util.spec_from_file_location(package + '.nodes.image_list', root / 'nodes/image_list.py')
node_module = importlib.util.module_from_spec(spec)
with patch.dict(sys.modules, modules):
spec.loader.exec_module(node_module)
start = node_module.DistributedImageListStart
start.hidden = SimpleNamespace(prompt=self.prompt, unique_id='start')
result = start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
context = result[3][0]
self.assertIsNotNone(runtime.get_master_job(context))
runtime.remove_master_job(context)
server.PromptServer.instance = self.server(['op', 'end'])
with self.assertRaisesRegex(ValueError, 'must be placed after'):
start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
self.assertEqual(runtime._MASTER_JOBS, {})
def test_owned_target_controlled_start_can_create_job(self):
context = runtime.create_master_job('start', 'end', self.prompt, [], [], '',
master_participates=True, explicit_targets=True)
self.addCleanup(runtime.remove_master_job, context)
self.assertIsNotNone(runtime.get_master_job(context))
class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
with runtime._LOCK:
runtime._MASTER_JOBS.clear()
runtime._WORKER_ITEMS.clear()
def tearDown(self):
with runtime._LOCK:
runtime._MASTER_JOBS.clear()
runtime._WORKER_ITEMS.clear()
@staticmethod
def _state(prompt_id=None, *, active=False, age=10_000.0):
return SimpleNamespace(
prompt_id=prompt_id,
active=active,
created_at=time.monotonic() - age,
)
@staticmethod
def _server(queue):
return SimpleNamespace(app=web.Application(), prompt_queue=queue)
async def _start(self, server, *, interval=0.01, grace=0.01):
lifecycle.install_lifecycle(
server,
runtime,
interval_seconds=interval,
orphan_grace_seconds=grace,
)
runner = web.AppRunner(server.app)
await runner.setup()
return runner
async def test_unattended_sweep_removes_orphan_and_expired_staging(self):
context_id = "orphan"
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)[context_id] = self._state()
runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"})
queue = SimpleNamespace(
get_current_queue_volatile=lambda: ([], []),
)
server = self._server(queue)
swept = asyncio.Event()
sweep_worker_items = runtime._sweep_worker_items
def track_worker_sweep():
swept.set()
sweep_worker_items()
with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep):
runner = await self._start(server)
try:
await asyncio.wait_for(swept.wait(), timeout=0.5)
self.assertNotIn(context_id, runtime._MASTER_JOBS)
self.assertNotIn("expired-token", runtime._WORKER_ITEMS)
with self.assertRaisesRegex(ValueError, "Unknown or expired"):
runtime.set_master_result(context_id, 0, "late-token", {})
finally:
await runner.cleanup()
async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self):
states = {
"active": self._state("active-prompt", active=True),
"queued": self._state("queued-prompt"),
"orphan": self._state("missing-prompt"),
"unbound": self._state(None),
}
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS).update(states)
queue = SimpleNamespace(
get_current_queue_volatile=lambda: (
[[1, "active-prompt", {}]],
[(2, "queued-prompt", {}, {}, None, None)],
),
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"])
self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"])
self.assertNotIn("orphan", runtime._MASTER_JOBS)
self.assertNotIn("unbound", runtime._MASTER_JOBS)
finally:
await runner.cleanup()
async def test_falls_back_to_nonvolatile_queue_reader(self):
state = self._state("queued-prompt")
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["queued"] = state
fallback_reads = 0
def volatile_failure():
raise RuntimeError("volatile view unavailable")
def current_queue():
nonlocal fallback_reads
fallback_reads += 1
return ([], [[1, "queued-prompt", {}]])
queue = SimpleNamespace(
get_current_queue_volatile=volatile_failure,
get_current_queue=current_queue,
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertGreaterEqual(fallback_reads, 1)
self.assertIs(runtime._MASTER_JOBS.get("queued"), state)
finally:
await runner.cleanup()
async def test_queue_read_error_protects_bound_contexts(self):
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt")
cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None)
def fail_read():
raise RuntimeError("queue unavailable")
queue = SimpleNamespace(
get_current_queue_volatile=fail_read,
get_current_queue=fail_read,
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertIn("bound", runtime._MASTER_JOBS)
self.assertNotIn("unbound", runtime._MASTER_JOBS)
finally:
await runner.cleanup()
async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self):
queue_reads = 0
def read_queue():
nonlocal queue_reads
queue_reads += 1
return ([], [])
server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue))
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
self.assertEqual(len(server.app.cleanup_ctx), 1)
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["removed"] = self._state()
runtime.remove_master_job("removed")
runtime.remove_master_job("removed")
self.assertNotIn("removed", runtime._MASTER_JOBS)
runner = web.AppRunner(server.app)
await runner.setup()
await asyncio.sleep(0.035)
self.assertGreaterEqual(queue_reads, 1)
await runner.cleanup()
reads_after_cleanup = queue_reads
await asyncio.sleep(0.035)
self.assertEqual(queue_reads, reads_after_cleanup)
if __name__ == "__main__":
unittest.main()
+437
View File
@@ -0,0 +1,437 @@
"""Master slot regressions, plus an opt-in real CPU ComfyUI subprocess probe."""
import asyncio
import importlib
import json
import os
from pathlib import Path
import subprocess
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
import torch
ROOT = Path(__file__).resolve().parents[1]
def graph():
return {
'source': {'class_type': 'Source', 'inputs': {}},
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['op', 0], 'context': ['start', 3]}},
}
class Response:
status = 200
def __init__(self, data=None):
self.data = data
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return None
async def json(self):
return self.data
def raise_for_status(self):
pass
class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
name = '_master_slot_tests'
self.modules = patch.dict(sys.modules)
self.modules.start()
self.addCleanup(self.modules.stop)
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
sys.modules[name + suffix] = module
self.rt = importlib.import_module(name + '.image_list.runtime')
self.workers = []
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
self.dispatched = []
self.local_calls = []
self.tokens = []
self.closed = 0
self.local_error = None
self.local_delay = 0.005
self.interrupted = False
self.local_started = asyncio.Event()
self.local_finished = asyncio.Event()
self.deadline = 0
self.slot_instances = []
self.config = {'workers': self.workers, 'settings': {}}
fixture = self
class FakeSlot:
def __init__(self, model_callback=None):
fixture.slot_instances.append(self)
self.tasks = []
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
async def execute():
inputs = prompt[end_id]['inputs']
index = inputs['worker_item_index']
token = inputs['worker_result_token']
fixture.local_calls.append(index)
fixture.tokens.append(token)
fixture.local_started.set()
try:
await asyncio.sleep(fixture.local_delay)
if fixture.local_error:
raise fixture.local_error
item = fixture.rt.consume_worker_item(token)
fixture.rt.set_master_result(fixture.context, index, token, {
'image': item['image'], 'mask': item.get('mask'),
'text': item['values'].get('value_1')})
finally:
fixture.local_finished.set()
task = asyncio.create_task(execute())
self.tasks.append(task)
return task
def request_stop(self):
pass
async def close(self):
# Model a synchronous node that must finish before its thread can join.
await asyncio.gather(*self.tasks, return_exceptions=True)
fixture.closed += 1
dependencies = {
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=FakeSlot),
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=self.preflight,
dispatch_worker_prompt=AsyncMock(side_effect=self.dispatch)),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda config, ids: self.workers),
'.api.orchestration.media_sync': types.SimpleNamespace(
sync_worker_media=AsyncMock(), fetch_worker_path_separator=AsyncMock(return_value=None)),
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
get_worker_timeout_seconds=lambda: 60,
get_image_list_job_timeout_seconds=lambda: self.deadline),
'.utils.network': types.SimpleNamespace(
build_master_callback_url=lambda *a, **kw: 'http://master',
build_worker_url=lambda worker, path: 'http://remote' + path,
get_client_session=AsyncMock(return_value=self)),
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
}
for suffix, module in dependencies.items():
sys.modules[name + suffix] = module
self.network = dependencies['.utils.network']
sys.modules['aiohttp'] = types.SimpleNamespace(ClientTimeout=lambda **kw: kw, ClientError=ConnectionError)
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
sys.modules['comfy'] = comfy
sys.modules['comfy.model_management'] = comfy.model_management
def create(self, participates=True, count=5):
self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)},
'index': i}
for i in range(count)]
self.context = self.rt.create_master_job('start', 'end', graph(), self.items,
[w['id'] for w in self.workers], '', master_participates=participates)
self.state = self.rt.get_master_job(self.context)
return self.context
def post(self, url, *, data=None, **kwargs):
if url.endswith('/distributed/image-list/item'):
payload = json.loads(data)
self.rt.store_worker_item(payload['token'], payload['item'])
return Response({})
def delete(self, url, **kwargs):
self.rt.delete_worker_item(url.rsplit('/', 1)[-1])
return Response({})
def get(self, url, **kwargs):
pid = url.rsplit('/', 1)[-1]
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
async def dispatch(self, worker, prompt, meta, **kwargs):
inputs = prompt['end']['inputs']
index, token = inputs['worker_item_index'], inputs['worker_result_token']
self.dispatched.append(index)
# A remote callback is ready immediately; the local slot is slower.
item = self.rt.consume_worker_item(token)
self.rt.set_master_result(self.context, index, token, {
'image': item['image'], 'mask': item['mask'], 'text': item['values']['value_1']})
return f'remote-{index}'
async def test_local_only_preserves_order_and_uses_no_http(self):
result = await self.rt.run_image_list_job(self.create())
self.assertEqual(result[2], [str(i) for i in range(5)])
self.assertEqual(self.local_calls, list(range(5)))
self.assertEqual(len(self.slot_instances), 1)
self.assertEqual(len(set(self.tokens)), 5)
self.network.get_client_session.assert_not_awaited()
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
for i, (image, mask) in enumerate(zip(*result[:2])):
self.assertTrue(torch.equal(image, self.items[i]['image']))
self.assertTrue(torch.equal(mask, self.items[i]['mask']))
async def test_mixed_pool_claims_from_one_queue_and_relaxes_remote_idle(self):
self.workers.append({'id': 'master', 'type': 'local'}) # Not the in-process master slot.
self.local_delay = 0.04
result = await self.rt.run_image_list_job(self.create(count=9))
self.assertTrue(self.local_calls)
self.assertTrue(self.dispatched)
self.assertEqual(sorted(self.local_calls + self.dispatched), list(range(9)))
self.assertEqual(result[2], [str(i) for i in range(9)])
self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
async def test_orchestrator_mode_never_instantiates_local_slot(self):
self.workers.append({'id': 'worker', 'type': 'remote'})
await self.rt.run_image_list_job(self.create(participates=False))
self.assertFalse(self.slot_instances)
self.preflight.assert_awaited_once_with(self.workers, require_idle=True)
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
self.local_error = RuntimeError('local branch failed')
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
await self.rt.run_image_list_job(self.create())
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_deadline_joins_local_execution_before_removing_context(self):
self.deadline, self.local_delay = 0.03, 0.1
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
await self.rt.run_image_list_job(self.create())
self.assertTrue(self.local_finished.is_set())
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_cancelled_scheduler_joins_owned_execution_even_if_cancelled_again(self):
self.local_delay = 0.12
task = asyncio.create_task(self.rt.run_image_list_job(self.create()))
await self.local_started.wait()
task.cancel()
await asyncio.sleep(0.02)
task.cancel()
with self.assertRaises(self.rt.ImageListCancelledError):
await task
self.assertTrue(self.local_finished.is_set())
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
class RealFrameworkTests(unittest.TestCase):
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for the real CPU framework probe')
def test_real_framework_local_slot(self):
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--real-probe',
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=60)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn('MASTER_SLOT_REAL_OK', result.stdout)
def real_framework_probe(comfy_root):
import threading
import time
sys.path.insert(0, str(Path(comfy_root).resolve()))
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import execution
import nodes
import comfy.memory_management as memory
import comfy.model_management as models
import comfy.utils as utils
from comfy_execution import progress
from comfy_execution.graph_utils import GraphBuilder
from comfy_api.v0_0_2 import io
name = '_real_master_slot'
for suffix in ('', '.image_list', '.utils', '.nodes'):
module = types.ModuleType(name + suffix)
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
sys.modules[name + suffix] = module
# No HTTP/service/extension bootstrap is needed by the local slot. Fail if
# either boundary tries to use a server-loop bridge instead of direct memory.
def forbidden(*args, **kwargs):
raise AssertionError('local boundary used the server loop or /prompt')
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
config = {'settings': {}}
sys.modules[name + '.utils.config'] = types.SimpleNamespace(load_config=lambda: config)
rt = importlib.import_module(name + '.image_list.runtime')
local = importlib.import_module(name + '.image_list.local_executor')
boundary = importlib.import_module(name + '.nodes.image_list')
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
DistributedImageListEnd=boundary.DistributedImageListEnd)
loads, calls = [], []
entered, stopped = threading.Event(), threading.Event()
mode = ['normal']
model = object()
class Loader:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('MODEL',)
FUNCTION = 'run'
def run(self):
loads.append(True)
return (model,)
class Operation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'model': ('MODEL',),
'mask': ('MASK',), 'text': ('STRING', {'forceInput': True})}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
async def run(self, image, model, mask, text):
calls.append(text)
entered.set()
if mode[0] == 'error':
raise ValueError('probe branch failure')
if mode[0] in ('cancel', 'interrupt'):
try:
while True:
utils.ProgressBar(1).update(1)
await asyncio.sleep(0.01)
finally:
await asyncio.sleep(0.02)
stopped.set()
if mode[0] == 'sync':
# No progress/interrupt checks: cleanup still must join, not
# release the parent just because an asyncio Future is cancelled.
time.sleep(0.08)
stopped.set()
await asyncio.sleep(0)
return image + 1, mask + 1, text + ':done'
class Expanded:
@classmethod
def INPUT_TYPES(cls):
return Operation.INPUT_TYPES()
RETURN_TYPES = Operation.RETURN_TYPES
FUNCTION = 'run'
def run(self, image, model, mask, text):
builder = GraphBuilder()
node = builder.node('Operation', image=image, model=model, mask=mask, text=text)
return {'result': (node.out(0), node.out(1), node.out(2)), 'expand': builder.finalize()}
nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded)
prompt = graph()
prompt['loader'] = {'class_type': 'Loader', 'inputs': {}}
prompt['op']['inputs'].update(model=['loader', 0], mask=['start', rt.MASK_OUT],
text=['start', rt.VALUE_OUT_BASE])
prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2])
items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, 'index': i}
for i in range(3)]
region = rt.validate_image_list_regions(prompt)[0]
parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({}))
parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7}
progress.global_progress_registry = parent_progress
parent_hook = lambda *a, **kw: None
utils.PROGRESS_BAR_HOOK = parent_hook
ram_callback = lambda *a, **kw: None
memory.set_ram_cache_release_state(ram_callback, 123)
GraphBuilder.set_default_prefix('outer-prefix', 7, 9)
def restored():
assert progress.global_progress_registry is parent_progress
assert parent_progress.nodes == {'outer': {'state': progress.NodeState.Running, 'value': 3, 'max': 7}}
assert utils.PROGRESS_BAR_HOOK is parent_hook
assert memory.extra_ram_release_callback is ram_callback and memory.RAM_CACHE_HEADROOM == 123
assert (GraphBuilder._default_prefix_root, GraphBuilder._default_prefix_call_index,
GraphBuilder._default_prefix_graph_index) == ('outer-prefix', 7, 9)
assert not any(t.name.startswith('image-list-master') for t in threading.enumerate())
async def probe():
# Verify actual V3 input parsing and list mapping, not only direct execute.
for supplied, delegated, expected in [(None, False, True), (None, True, False),
(False, False, False), (True, True, True)]:
config['settings']['master_delegate_only'] = delegated
inputs = {'image_list_enabled_worker_ids': '["remote"]'}
if supplied is not None:
inputs['image_list_master_participates'] = supplied
cls = boundary.DistributedImageListStart
data, _, v3 = execution.get_input_data(inputs, cls, 'start', dynprompt=execution.DynamicPrompt(prompt))
data['images'] = [torch.zeros((3, 2, 2, 3))]
result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3)
result = await execution.resolve_map_node_over_list_results(result)
context = result[0].result[rt.CONTEXT_OUT][0]
assert rt.get_master_job(context).master_participates is expected
rt.remove_master_job(context)
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
mode[0] = behavior
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
loads.clear()
calls.clear()
entered.clear()
stopped.clear()
tracked = []
slot = local.LocalExecutionSlot(model_callback=tracked.append)
context = rt.create_master_job('start', 'end', prompt, items, [], '', master_participates=True)
try:
for index in range(3 if behavior in ('normal', 'expand') else 1):
token = rt.create_worker_task_token()
rt.register_assignment(context, index, 'local', token)
rt._WORKER_ITEMS[token] = (time.monotonic() + 60, items[index])
worker_prompt = rt.build_worker_prompt(prompt, region, token, context, index, '')
for key in ('start', 'end'):
worker_prompt[key]['inputs']['image_list_local_execution'] = True
if behavior == 'already-interrupted':
models.interrupt_current_processing(True)
future = slot.submit(worker_prompt, f'local-{index}', 'end')
if behavior in ('cancel', 'interrupt', 'sync'):
assert await asyncio.to_thread(entered.wait, 5)
if behavior == 'interrupt':
models.interrupt_current_processing(True)
else:
slot.request_stop()
try:
await asyncio.wait_for(asyncio.shield(asyncio.wrap_future(future)), 5)
except rt.ImageListCancelledError:
assert behavior in ('cancel', 'interrupt', 'sync', 'already-interrupted'), behavior
except RuntimeError as exc:
assert behavior == 'error' and 'probe branch failure' in str(exc), str(exc)
else:
assert behavior in ('normal', 'expand'), behavior
result = rt.get_master_job(context).results[index]
assert torch.equal(result['image'], items[index]['image'] + 1)
assert torch.equal(result['mask'], items[index]['mask'] + 1)
assert result['text'] == f'{index}:done'
if behavior in ('normal', 'expand'):
assert len(loads) == 1 and calls == ['0', '1', '2'], (loads, calls)
assert any(model in outputs for outputs in tracked), tracked
if behavior in ('cancel', 'interrupt', 'sync'):
assert stopped.is_set(), behavior
if behavior == 'already-interrupted':
assert models.processing_interrupted() and not calls
finally:
await slot.close()
for token in rt.get_master_job(context).expected_tokens.values():
rt.delete_worker_item(token)
rt.remove_master_job(context)
models.interrupt_current_processing(False) # Only test setup/teardown.
restored()
assert not rt._MASTER_JOBS and not rt._WORKER_ITEMS
asyncio.run(probe())
print('MASTER_SLOT_REAL_OK V3 metadata; cache reuse; async/expanded graph; errors; cancellation; state restoration')
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == '--real-probe':
real_framework_probe(sys.argv[2])
else:
unittest.main()
+147
View File
@@ -0,0 +1,147 @@
"""Image-list participation uses the existing master orchestrator setting."""
import json
import sys
import unittest
from unittest.mock import AsyncMock, patch
from test_serverless_integration import PKG, load_modules
class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.q, modules, self.ps = load_modules()
self.modules_patch = patch.dict(sys.modules, modules)
self.modules_patch.start()
self.addCleanup(self.modules_patch.stop)
self.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
self.config = {"workers": [self.worker], "settings": {}}
self.q.load_config = lambda: self.config
# Model the legacy helper's fallback: image-list must not inherit it.
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
(workers, delegate if workers else False))
self.preflight = AsyncMock(return_value={"worker": True})
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
self.prompt = {
"1": {"class_type": "EmptyImage", "inputs": {}},
"2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
"3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}},
}
async def queue(self, **kwargs):
return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)
def prepared_inputs(self):
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
async def test_prepared_master_prunes_side_outputs_without_mutating_input(self):
import copy
self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}}
self.prompt['3']['inputs']['images'] = ['op', 0]
self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}}
self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}}
original = copy.deepcopy(self.prompt)
registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': {
'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}})
with patch.dict(sys.modules, {'nodes': registry}):
await self.queue()
queued = self.q.queue_prompt_payload.call_args
self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save'])
self.assertIn('op', queued.args[0])
self.assertNotIn('preview', queued.args[0])
self.assertEqual(self.prompt, original)
async def test_default_includes_master_and_direct_workers(self):
result = await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
self.assertEqual(result[2], 1) # Public worker_count still counts remote participants only.
self.preflight.assert_awaited_once_with(
[{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
require_idle=False)
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_orchestrator_setting_excludes_master_processing(self):
self.config["settings"]["master_delegate_only"] = True
await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
# Upstream source stays on master even in orchestrator-only mode.
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
async def test_explicit_override_wins_over_config_in_both_directions(self):
for configured, requested in [(True, False), (False, True)]:
with self.subTest(configured=configured, requested=requested):
self.config["settings"]["master_delegate_only"] = configured
await self.queue(delegate_master=requested)
self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)
async def test_master_only_with_explicit_empty_worker_selection(self):
result = await self.queue(enabled_worker_ids=[])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertEqual(result[2], 0)
self.preflight.assert_not_awaited()
async def test_master_only_when_no_workers_are_configured(self):
self.config["workers"] = []
await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
async def test_unavailable_workers_do_not_prevent_enabled_master(self):
self.q.select_active_workers.return_value = ([], False)
self.q.select_active_workers.side_effect = None
self.preflight.return_value = {}
await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
async def test_busy_workers_do_not_prevent_enabled_master(self):
self.preflight.return_value = {}
await self.queue()
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
async def test_orchestrator_without_workers_fails_before_queue(self):
self.config["settings"]["master_delegate_only"] = True
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
await self.queue(enabled_worker_ids=[])
self.q.queue_prompt_payload.assert_not_awaited()
async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
self.config["settings"]["master_delegate_only"] = True
self.q.select_active_workers.side_effect = None
self.q.select_active_workers.return_value = ([], False)
self.preflight.return_value = {}
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_serverless_pool_is_retained_with_participating_master(self):
self.config["workers"] = [{
"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"endpoint_id": "endpoint", "media_profile": "comfy-media-v1",
}]
result = await self.queue()
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["pool"])
self.assertEqual(result[2], 1)
self.preflight.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_legacy_serverless_pool_is_rejected_not_silently_ignored(self):
self.config["workers"] = [{
"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"endpoint_id": "endpoint",
}]
with self.assertRaisesRegex(ValueError, "comfy-media-v1"):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
self.preflight.assert_not_awaited()
async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
self.preflight.side_effect = RuntimeError("upgrade worker")
with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
+358
View File
@@ -0,0 +1,358 @@
"""Exercise production per-item orchestration through controlled provider seams."""
import asyncio
import importlib.util
import json
from pathlib import Path
import sys
import types
import unittest
from unittest.mock import AsyncMock, Mock, patch
import torch
from test_image_list import graph_one_region
ROOT = Path(__file__).resolve().parents[1]
class Response:
def __init__(self, data=None, status=200):
self.data, self.status = data, status
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.data
def raise_for_status(self):
if self.status >= 400:
raise RuntimeError(f'HTTP {self.status}')
class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
name = '_image_list_boundaries'
self.modules = {}
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = []
self.modules[name + suffix] = module
self.patch_modules = patch.dict(sys.modules, self.modules)
self.patch_modules.start()
self.addCleanup(self.patch_modules.stop)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
self.rt = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = self.rt
spec.loader.exec_module(self.rt)
self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188}
self.config = {'workers': [self.worker], 'settings': {}}
self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], []
self.reject = False
self.interrupted = False
self.completed = True
self.queue = {'queue_running': [], 'queue_pending': []}
self.callback = Mock(side_effect=lambda worker, **_kwargs:
'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master')
self.sync = AsyncMock(side_effect=self.sync_media)
self.dispatch = AsyncMock(side_effect=self.dispatch_prompt)
self.logs = []
dependencies = {
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=AsyncMock(return_value={'worker': True}),
dispatch_worker_prompt=self.dispatch),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda config, ids: config['workers']),
'.api.orchestration.media_sync': types.SimpleNamespace(
sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'),
convert_paths_for_platform=lambda prompt, separator: prompt),
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
get_worker_timeout_seconds=lambda: 60,
get_image_list_job_timeout_seconds=lambda config=None: 0),
'.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master',
build_master_callback_url=self.callback,
build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}",
get_client_session=AsyncMock(return_value=self)),
'.utils.logging': types.SimpleNamespace(log=self.logs.append),
}
for suffix, module in dependencies.items():
sys.modules[name + suffix] = module
self.config_module = dependencies['.utils.config']
# Other legacy test modules replace aiohttp during collection. This
# controlled provider fixture needs only its explicit timeout/error seam.
sys.modules['aiohttp'] = types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError)
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
sys.modules['comfy'] = comfy
sys.modules['comfy.model_management'] = comfy.model_management
async def sync_media(self, worker, prompt):
prompt['7']['inputs']['image'] = 'shared/ref.png'
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
self.assertEqual(kwargs.get('partial_execution_targets'), ['5'])
if self.reject:
raise RuntimeError('prompt rejected')
self.prompts.append((json.loads(json.dumps(prompt)), meta))
inputs = prompt['5']['inputs']
token = inputs['worker_result_token']
item = self.rt.consume_worker_item(token)
self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']})
return f'prompt-{len(self.prompts)}'
def post(self, url, *, json=None, data=None, **kwargs):
if data is not None:
json = __import__('json').loads(data)
if url.endswith('/distributed/image-list/item'):
self.uploads.append(json['token'])
self.rt.store_worker_item(json['token'], json['item'])
else:
self.cancel_requests.append((url, json))
return Response({})
def delete(self, url, **kwargs):
token = url.rsplit('/', 1)[-1]
self.deletions.append(token)
self.rt.delete_worker_item(token)
return Response({})
def get(self, url, **kwargs):
if url.endswith('/queue'):
return Response(self.queue)
pid = url.rsplit('/', 1)[-1]
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}}
if self.completed else {})
def create(self, count=2):
prompt = graph_one_region()
prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}}
prompt['4']['inputs']['reference'] = ['7', 0]
self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}}
self.context = self.rt.create_master_job('3', '5', prompt,
[{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)],
['worker'], 'http://public-master', self.metadata)
return self.context
async def test_runtime_preflight_rejects_old_target_capability_before_upload(self):
from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability
real_dispatch = _load_dispatch_module()
capability = _valid_capability()
capability.pop('execution_targets')
session = _FakeSession({('worker', 'capabilities'): (200, capability)})
module = sys.modules['_image_list_boundaries.api.orchestration.dispatch']
module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers
with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object(
real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'):
await self.rt.run_image_list_job(self.create(1))
self.assertFalse(self.uploads)
self.dispatch.assert_not_awaited()
self.assertFalse(self.rt._MASTER_JOBS)
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
context = self.create()
await self.rt.run_image_list_job(context)
self.assertEqual(self.sync.await_count, 1)
self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts],
['shared/ref.png', 'shared/ref.png'])
tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts]
self.assertEqual(len(set(tokens)), 2)
self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1])
self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata])
async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self):
import tempfile
context = self.create(2)
state = self.rt.get_master_job(context)
state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}}
state.original_prompt['4']['inputs']['second_reference'] = ['8', 0]
name = '_image_list_boundaries.api.orchestration.media_sync'
spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py')
assert spec is not None and spec.loader is not None
media = importlib.util.module_from_spec(spec)
sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append
sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace(
trace_debug=lambda *args: None, trace_info=lambda *args: None)
spec.loader.exec_module(media)
media.fetch_worker_path_separator = AsyncMock(return_value='\\')
converted = Mock(wraps=media.convert_paths_for_platform)
media.convert_paths_for_platform = converted
uploads = []
async def upload(worker, filename, content, digest, mime, **kwargs):
uploads.append((filename, content))
return True, f'worker-input/{len(uploads)}.png'
media._upload_media_to_worker = AsyncMock(side_effect=upload)
media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media)
with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory:
root = Path(directory)
for kind in ('input', 'output', 'temp'):
path = root / kind / 'shared/ref.png'
path.parent.mkdir(parents=True)
path.write_bytes(kind.encode())
def resolve_annotated(filename):
for kind in ('output', 'temp', 'input'):
suffix = f' [{kind}]'
if filename.endswith(suffix):
return str(root / kind / filename[:-len(suffix)])
return str(root / 'input' / filename)
with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace(
get_annotated_filepath=resolve_annotated)}):
await self.rt.run_image_list_job(context)
self.assertEqual(uploads, [('shared/ref.png [output]', b'output'),
('shared/ref.png [temp]', b'temp')])
media.sync_worker_media.assert_awaited_once()
media.fetch_worker_path_separator.assert_awaited_once()
converted.assert_called_once()
self.assertEqual(converted.call_args.args[1], '\\')
for prompt, _ in self.prompts:
self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png')
self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png')
self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]')
async def test_local_worker_callback_does_not_use_public_master_address(self):
self.worker['type'] = 'local'
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188')
self.callback.assert_called()
async def test_remote_worker_uses_topology_callback_helper(self):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master')
self.callback.assert_called()
def test_metadata_is_forwarded_into_master_start_prompt(self):
from test_prompt_transform import pt
prompt = graph_one_region()
metadata = {'nodes': [{'id': 1}]}
index = pt.PromptIndex(prompt)
transformed = pt.apply_participant_overrides(
prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False,
index, workflow_meta=metadata)
self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata)
async def test_aggregate_input_limit_fails_before_upload(self):
self.rt.MAX_REQUEST_BYTES = 64
with self.assertRaisesRegex(ValueError, 'complete.*request'):
await self.rt.run_image_list_job(self.create(1))
self.assertFalse(self.uploads)
self.assertFalse(self.rt._WORKER_ITEMS)
def test_tensor_and_combined_request_limits_preserve_dtype(self):
item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0}
encoded = self.rt.encode_item(item)
payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded}
single = self.rt.encode_request_payload(payload)
self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32')
self.rt.MAX_REQUEST_BYTES = len(single)
self.assertEqual(self.rt.encode_request_payload(payload), single)
item['values'] = {'value_1': item['image']}
item['mask'] = torch.zeros((1, 4, 4))
payload['item'] = self.rt.encode_item(item)
with self.assertRaisesRegex(ValueError, 'complete.*request'):
self.rt.encode_request_payload(payload)
# Size is checked before allocating/copying a 4K tensor onto the CPU.
with self.assertRaisesRegex(ValueError, 'transport limit'):
self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE')
async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self):
self.reject = True
with self.assertRaisesRegex(RuntimeError, 'prompt rejected'):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertFalse(self.cancel_requests)
async def test_user_interrupt_during_pending_submission_is_bounded(self):
entered = asyncio.Event()
async def never_ack(*args, **kwargs):
entered.set()
await asyncio.Event().wait()
self.dispatch.side_effect = never_ack
self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02
task = asyncio.create_task(self.rt.run_image_list_job(self.create(2)))
await asyncio.wait_for(entered.wait(), 1)
self.interrupted = True
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 0.8)
self.assertEqual(len(self.uploads), 1)
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs)
async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self):
self.completed = False
context = self.create(2)
clock = [0.0]
calls = [0]
original_get = self.get
def get(url, **kwargs):
if '/history/' in url:
calls[0] += 1
clock[0] += 3700
self.completed = calls[0] >= 3
if url.endswith('/queue'):
return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [],
'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []})
return original_get(url, **kwargs)
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
with patch.object(self, 'get', side_effect=get):
result = await self.rt.run_image_list_job(context)
self.assertEqual(len(result[0]), 2)
self.assertGreaterEqual(clock[0], 7200)
self.assertFalse(self.cancel_requests)
async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self):
self.completed = False
clock = [0.0]
original_get = self.get
def get(url, **kwargs):
if '/history/' in url:
clock[0] += 70
if url.endswith('/queue'):
return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []})
return original_get(url, **kwargs)
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
with patch.object(self, 'get', side_effect=get):
with self.assertRaisesRegex(TimeoutError, 'liveness'):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.cancel_requests, [
('http://worker/queue', {'delete': ['prompt-1']}),
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self):
self.completed = False
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
original_post = self.post
def post(url, **kwargs):
result = original_post(url, **kwargs)
return Response(status=500) if url.endswith('/queue') else result
with patch.object(self, 'post', side_effect=post):
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
await self.rt.run_image_list_job(self.create(2))
self.assertEqual(self.cancel_requests, [
('http://worker/queue', {'delete': ['prompt-1']}),
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
self.assertEqual(self.deletions, self.uploads)
self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs)
self.assertEqual(self.dispatch.await_count, 1)
async def test_configured_total_deadline_cleans_owned_work(self):
self.completed = False
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []}
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1)
self.assertEqual(len(self.uploads), 1)
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._MASTER_JOBS)
+683
View File
@@ -0,0 +1,683 @@
"""CPU-only image-list compilation and lifecycle tests; providers are mocked."""
import asyncio
import base64
import copy
import importlib
import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
import pytest
import torch
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "_image_list_serverless_tests"
for suffix in ("", ".api", ".api.orchestration", ".image_list", ".serverless_bridge"):
module = types.ModuleType(PACKAGE + suffix)
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[module.__name__] = module
adapter = importlib.import_module(PACKAGE + ".image_list.serverless")
runtime = adapter.runtime
manager_module = importlib.import_module(PACKAGE + ".api.orchestration.serverless_runtime")
handler = importlib.import_module(PACKAGE + ".serverless_bridge.handler")
codec, contract = adapter.codec, adapter.contract
def node(kind, **inputs):
return {"class_type": kind, "inputs": inputs}
def pool(wid="pool", jobs=1, enabled=True, **overrides):
return dict(id=wid, type="serverless", endpoint_id="ep", timeout_seconds=30,
jobs_per_run=jobs, enabled=enabled, media_profile=contract.PROFILE, **overrides)
def source():
return {
"upstream": node("LoadImage", image="never-read.png"),
"ns/start": node(runtime.START_CLASS, images=["upstream", 0]),
"model": node("CheckpointLoaderSimple", ckpt_name="fixed.safetensors"),
"process": node("Process", image=["ns/start", 0], mask=["ns/start", 1],
index=["ns/start", 2], text=["ns/start", 4], integer=["ns/start", 5],
floating=["ns/start", 6], boolean=["ns/start", 7], model=["model", 0]),
"ns/end": node(runtime.END_CLASS, context=["ns/start", 3], images=["process", 0],
masks=["process", 1], texts=["ns/start", 4], worker_result_token="drop-me"),
"downstream": node("SaveImage", images=["ns/end", 0]),
}
def region(prompt=None):
return {"start_id": "ns/start", "end_id": "ns/end",
"node_ids": set(prompt or source()) - {"upstream", "downstream"}}
def item(index=0, dtype=torch.float32):
return {"index": index, "image": torch.arange(24, dtype=dtype).reshape(1, 2, 3, 4) / 24,
"mask": torch.arange(6, dtype=dtype).reshape(1, 2, 3) / 6,
"values": {"value_1": " \n雪", "value_2": 123, "value_3": 0.25, "value_4": False}}
async def prepare(prompt=None, value=None, worker=None):
prompt = prompt or source()
return await adapter.prepare_serverless_item(prompt, region(prompt), value or item(), worker or pool(),
request_id="request:1", job_id="queue:1")
async def decoded_inputs(job):
return {artifact["id"]: codec.decode_payload(artifact["kind"], artifact["descriptor"],
await contract.resolve_artifact(artifact)) for artifact in job["media"]["artifacts"]}
def test_selection_is_exact_no_expansion_or_profile_fallback():
a, b = pool("a", jobs=16), pool("b", enabled=False)
direct = {"id": "direct", "type": "remote", "enabled": True}
config = {"workers": [a, b, direct]}
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
resolved = adapter.resolve_serverless_workers(config, ["b", "direct"])
assert [w["id"] for w in resolved] == ["b"]
assert resolved[0]["pool_id"] == "b"
resolved[0]["jobs_per_run"] = 2
assert b["jobs_per_run"] == 1
assert adapter.resolve_serverless_workers(config, []) == []
legacy = pool("legacy", enabled=False)
legacy.pop("media_profile")
config["workers"].append(legacy)
assert [w["id"] for w in adapter.resolve_serverless_workers(config, None)] == ["a"]
with pytest.raises(ValueError, match="comfy-media-v1"):
adapter.resolve_serverless_workers(config, ["legacy"])
with pytest.raises(ValueError, match="Unknown"):
adapter.resolve_serverless_workers(config, ["a:1"])
with pytest.raises(ValueError, match="Invalid"):
adapter.resolve_serverless_workers(config, [" a"])
with pytest.raises(ValueError, match="Duplicate"):
adapter.resolve_serverless_workers({"workers": [a, a]}, ["a"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_deterministic_graph_exact_codecs_and_source_immutable(dtype):
async def check():
prompt, value = source(), item(dtype=dtype)
original, saved_image = copy.deepcopy(prompt), value["image"].clone()
one = await prepare(prompt, value)
two = await prepare(dict(reversed(list(prompt.items()))), value)
assert one == two
assert prompt == original and torch.equal(value["image"], saved_image)
assert "upstream" not in one["workflow"] and "downstream" not in one["workflow"]
assert one["workflow"]["model"] == prompt["model"]
assert not any("distributed" in n["class_type"].lower() for n in one["workflow"].values())
proc = one["workflow"]["process"]["inputs"]
for name, expected in [("index", 0), ("integer", 123), ("floating", 0.25), ("boolean", False), ("text", " \n雪")]:
assert proc[name] == expected and type(proc[name]) is type(expected)
output = next(b for b in one["manifest"] if b["socket"] == "texts")
text_source = one["workflow"][output["node_id"]]["inputs"]["texts"][0]
assert one["workflow"][text_source] == node("PrimitiveString", value=" \n雪")
values = await decoded_inputs(one)
for ref in one["media"]["inputs"]:
result = values[ref["artifact_id"]]
assert result.dtype == dtype
assert torch.equal(result, value[ref["input_name"]])
assert len(values) == 2
handler._validate_request({key: one[key] for key in ("workflow", "images", "media")})
asyncio.run(check())
def test_direct_passthrough_exports_and_tensor_generic_values_deduplicate():
async def check():
prompt, value = source(), item()
prompt["process"]["inputs"]["reference"] = ["ns/start", 8]
prompt["process"]["inputs"]["reference_mask"] = ["ns/start", 9]
value["values"].update(value_5=value["image"], value_6=value["mask"])
prompt["ns/end"]["inputs"]["images"] = ["ns/start", 0]
job = await prepare(prompt, value)
assert "process" in job["workflow"] # mask still depends on processing
assert len(job["media"]["artifacts"]) == 2
refs = job["media"]["inputs"]
assert len(refs) == 5
image_binding = next(b for b in job["manifest"] if b["socket"] == "images")
assert any(r["node_id"] == image_binding["node_id"] and r["input_name"] == "images" for r in refs)
decoded = await decoded_inputs(job)
for ref in refs:
tensor = decoded[ref["artifact_id"]]
assert torch.equal(tensor, value["mask"] if "mask" in ref["input_name"] else value["image"])
asyncio.run(check())
@pytest.mark.parametrize("bad", [None, object(), torch.ones(2), torch.ones(2, 3)])
def test_connected_unsupported_generic_values_fail(bad):
value = item()
value["values"]["value_1"] = bad
with pytest.raises((ValueError, TypeError)):
asyncio.run(prepare(value=value))
@pytest.mark.parametrize("bad", [True, -1, "0", 0.0])
def test_start_slot_identity_not_coerced(bad):
prompt = source()
prompt["process"]["inputs"]["image"] = ["ns/start", bad]
with pytest.raises(ValueError, match="link"):
asyncio.run(prepare(prompt))
@pytest.mark.parametrize("bad", ["../bad", "ns//bad", "ns/./bad", "bad ", 42])
def test_node_identity_not_normalized_even_on_pruned_branch(bad):
prompt = source()
prompt[bad] = node("PrimitiveString", value="unused")
with pytest.raises(ValueError):
asyncio.run(prepare(prompt))
def test_placeholder_namespace_collision_and_empty_text():
async def check():
prompt, value = source(), item()
prompt["image-list:input:0"] = node("PrimitiveInt", value=5)
prompt["process"]["inputs"]["fixed"] = ["image-list:input:0", 0]
value["values"]["value_1"] = ""
job = await prepare(prompt, value)
assert job["workflow"]["image-list:input:0"] == prompt["image-list:input:0"]
assert job["workflow"]["process"]["inputs"]["image"] == ["image-list:input:0:1", 0]
assert any(n == node("PrimitiveString", value="") for n in job["workflow"].values())
asyncio.run(check())
def test_no_connected_none_and_payload_budget_before_submission():
value = item()
value["mask"] = None
with pytest.raises(ValueError, match="empty"):
asyncio.run(prepare(value=value))
client = importlib.import_module(PACKAGE + ".api.orchestration.runpod_client")
with patch.object(client, "MAX_REQUEST_BYTES", 10):
with pytest.raises(ValueError, match="Serialized"):
asyncio.run(prepare())
def test_fixed_file_inputs_are_prepared_on_private_copy_and_merge(tmp_path, monkeypatch):
input_dir = tmp_path / "input"
input_dir.mkdir()
(tmp_path / "output").mkdir()
(tmp_path / "temp").mkdir()
payload = b"bounded original file bytes"
(input_dir / "fixed.png").write_bytes(payload)
folders = types.SimpleNamespace(get_input_directory=lambda: str(input_dir),
get_output_directory=lambda: str(tmp_path / "output"), get_temp_directory=lambda: str(tmp_path / "temp"),
get_annotated_filepath=lambda name: str(input_dir / name))
monkeypatch.setitem(sys.modules, "folder_paths", folders)
prompt = source()
prompt["fixed"] = node("LoadImage", image="fixed.png")
prompt["process"]["inputs"]["fixed"] = ["fixed", 0]
original = copy.deepcopy(prompt)
job = asyncio.run(prepare(prompt))
assert prompt == original
assert len(job["images"]) == 1
assert base64.b64decode(job["images"][0]["image"]) == payload
assert job["workflow"]["fixed"]["inputs"]["image"] == job["images"][0]["name"]
assert len(job["media"]["artifacts"]) == 2
def test_merge_rejects_conflicting_artifact_ids_and_combined_caps():
async def check():
job = await prepare()
artifact = job["media"]["artifacts"][0]
assert adapter._merge_artifacts([artifact], [copy.deepcopy(artifact)]) == [artifact]
bad = copy.deepcopy(artifact)
bad["sha256"] = "0" * 64
with pytest.raises(ValueError, match="identity"):
adapter._merge_artifacts([artifact], [bad])
with patch.object(contract, "MAX_TOTAL_BYTES", 1):
with pytest.raises(ValueError, match="limits"):
adapter._merge_artifacts([artifact])
asyncio.run(check())
def test_handler_injects_exact_native_values_and_prunes_placeholders():
async def check():
value = item(dtype=torch.float64)
prompt = source()
prompt["ns/end"]["inputs"].update(images=["ns/start", 0], masks=["ns/start", 1])
job = await prepare(prompt, value)
workflow, media, legacy, targets = handler._validate_request(
{key: job[key] for key in ("workflow", "images", "media")})
buffers = {}
def write(payload):
name = f"owned-{len(buffers)}.bin"
buffers[name] = payload
return name
class Submitted(Exception):
pass
async def submit(native_graph):
assert all(n["class_type"] in {"MediaExportV1", "MediaImportV1", "PrimitiveString"}
for n in native_graph.values())
for binding in job["manifest"]:
link = native_graph[binding["node_id"]]["inputs"][binding["socket"]]
native = native_graph[link[0]]
if binding["kind"] == "STRING":
assert native == node("PrimitiveString", value=value["values"]["value_1"])
continue
assert native["class_type"] == "MediaImportV1"
metadata = json.loads(native["inputs"]["artifact_json"])
decoded = codec.decode_payload(metadata["kind"], metadata["descriptor"],
buffers[native["inputs"]["filename"]])
assert link[1] == {"IMAGE": 0, "MASK": 2}[binding["kind"]]
assert torch.equal(decoded, value["image" if binding["kind"] == "IMAGE" else "mask"])
# Stop at the mocked HTTP boundary; no endpoint/process/provider exists.
raise Submitted()
execution = types.SimpleNamespace(media=media, targets=targets, remaining=lambda: 30,
stage=types.SimpleNamespace(write=write), submit_and_wait=submit)
with pytest.raises(Submitted):
await handler._run(execution, workflow, legacy, None, None)
asyncio.run(check())
def test_fixed_typed_inputs_merge_with_start_refs_without_losing_consumers():
async def check():
prompt = source()
prompt["process"]["inputs"]["fixed_audio"] = 0
audio = {"waveform": torch.ones(1, 2, 3), "sample_rate": 48000}
descriptor, payload = codec.encode_payload("AUDIO", audio)
artifact = contract.artifact_from_payload("fixed-audio", "AUDIO", descriptor, payload)
fixed = AsyncMock(return_value=([], [artifact], [
{"node_id": "process", "input_name": "fixed_audio", "artifact_id": "fixed-audio"}]))
with patch.object(adapter, "prepare_media_inputs", fixed):
job = await prepare(prompt)
assert len(job["media"]["artifacts"]) == 3
assert len(job["media"]["inputs"]) == 3
assert any(ref["input_name"] == "fixed_audio" for ref in job["media"]["inputs"])
decoded = await decoded_inputs(job)
assert decoded["fixed-audio"]["sample_rate"] == 48000
assert torch.equal(decoded["fixed-audio"]["waveform"], audio["waveform"])
asyncio.run(check())
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.server = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
self.manager = manager_module.ServerlessManager(self.server,
types.SimpleNamespace(get=lambda _wid: "fake-key"), Path(self.temp.name) / "jobs.json", AsyncMock())
self.manager.poll_seconds = 0
self.state = runtime.ImageListJob("context", "ns/start", "ns/end", source(),
[item(i) for i in range(40)], ["pool"], "", prompt_id="master-prompt")
self.gates, self.jobs, self.mutations = {}, {}, {}
self.started = asyncio.Event()
async def submit(workflow, images, timeout, *, media):
jid = "provider-" + str(len(self.jobs))
self.jobs[jid] = dict(workflow=workflow, images=images, media=media)
self.gates[jid] = asyncio.Event()
self.started.set()
return jid
async def status(jid):
await self.gates[jid].wait()
output = self.output(self.jobs[jid])
if jid in self.mutations:
self.mutations[jid](output)
return {"status": "COMPLETED", "output": output}
self.api = types.SimpleNamespace(submit=AsyncMock(side_effect=submit), status=AsyncMock(side_effect=status),
cancel=AsyncMock(return_value={"status": "CANCELLED"}))
self.patches = [patch.object(manager_module, "RunpodClient", return_value=self.api),
patch.object(adapter, "get_manager", return_value=self.manager)]
for p in self.patches:
p.start()
self.event = asyncio.Event()
self.controllers = []
self.controller = await self.controller_for()
async def controller_for(self, workers=None, state=None):
ctl = await adapter.create_serverless_slots(state or self.state, region(),
{"workers": workers or [pool()]}, self.event)
self.controllers.append(ctl)
return ctl
async def asyncTearDown(self):
for ctl in self.controllers:
await ctl.close()
await self.manager.cancel()
for p in reversed(self.patches):
p.stop()
self.temp.cleanup()
def output(self, job):
media = job["media"]
output = dict(version=1, profile=contract.PROFILE, request_id=media["request_id"],
bindings=[], artifacts=[], items=[])
index = int(media["request_id"].rsplit(":", 1)[1])
data = item(index)
values = {"images": data["image"], "masks": data["mask"], "texts": str(index)}
for n, binding in enumerate(media["outputs"]):
descriptor, payload = codec.encode_payload(binding["kind"], values[binding["socket"]])
aid = f"result:{n}"
output["bindings"].append(dict(binding, count=1))
output["artifacts"].append(contract.artifact_from_payload(aid, binding["kind"], descriptor, payload))
output["items"].insert(0, dict(binding_id=binding["binding_id"], item_index=0, artifact_id=aid))
return {"media": output}
async def dispatch(self, index=0, ctl=None, slot=None):
ctl = ctl or self.controller
slot = slot or ctl.slots[index]
handle = await ctl.dispatch(slot, index)
return slot, handle
async def finish(self, handle):
# Submission has run by the time dispatch's shield completes, but yield
# explicitly rather than assuming provider status has already been entered.
while len(self.jobs) <= handle["index"]:
await asyncio.sleep(0)
jid = next(j for j, job in self.jobs.items()
if job["media"]["request_id"] == handle["execution_id"])
self.gates[jid].set()
await handle["task"]
async def test_no_pool_does_not_initialize_manager_or_credentials(self):
state = copy.copy(self.state)
state.worker_ids = []
with patch.object(adapter, "get_manager", side_effect=AssertionError("must stay lazy")):
ctl = await adapter.create_serverless_slots(state, region(), {"workers": [pool()]}, self.event)
assert ctl.slots == []
await ctl.close()
async def test_slots_are_bounded_not_jobs_per_run_and_cannot_collide(self):
assert len(self.controller.slots) == manager_module.PER_WORKER_JOB_LIMIT == 16
ctl = await self.controller_for([pool(jobs=16)])
assert ctl.slots == self.controller.slots
assert all(isinstance(slot, tuple) and len(slot) == 3 for slot in ctl.slots)
assert len(set(ctl.slots)) == 16 and "pool" not in ctl.slots
many = [pool(f"p{i}") for i in range(8)]
state = copy.copy(self.state)
state.worker_ids = [w["id"] for w in many]
ctl = await self.controller_for(many, state)
assert len(ctl.slots) == 64
assert all(sum(s[1] == w["id"] for s in ctl.slots) <= 16 for w in many)
async def test_reverse_completion_order_and_slot_reuse(self):
first, second = await self.dispatch(0), await self.dispatch(1)
inflight = dict([first, second])
assert self.controller.poll(inflight) is None
await self.finish(second[1])
await asyncio.sleep(0)
assert self.event.is_set()
slot, index, result = self.controller.poll(inflight)
assert (slot, index) == (second[0], 1) and result["text"] == "1"
assert torch.equal(result["image"], item(1)["image"])
assert second[1]["job_id"] not in self.server.distributed_pending_jobs
del inflight[slot]
third = await self.dispatch(2, slot=slot)
assert third[1]["execution_id"] != second[1]["execution_id"]
await self.finish(first[1])
assert self.controller.poll(inflight)[1] == 0
await self.finish(third[1])
assert self.controller.poll(dict([third]))[1] == 2
async def test_many_items_submit_once_each_through_bounded_reusable_slots(self):
ctl = await self.controller_for([pool(jobs=16)])
first_wave = asyncio.Event()
active = 0
peak = 0
original_submit = self.api.submit.side_effect
async def submit(*args, **kwargs):
jid = await original_submit(*args, **kwargs)
if len(self.jobs) == 16:
first_wave.set()
return jid
async def status(jid):
nonlocal active, peak
active += 1
peak = max(peak, active)
try:
await first_wave.wait()
await asyncio.sleep(0.001)
return {"status": "COMPLETED", "output": self.output(self.jobs[jid])}
finally:
active -= 1
self.api.submit.side_effect = submit
self.api.status.side_effect = status
async def next_result(inflight):
while True:
self.event.clear()
completed = ctl.poll(inflight)
if completed is not None:
return completed
await self.event.wait()
async def cleanup(*_):
pass
results = await asyncio.wait_for(runtime.dynamic_schedule(
len(self.state.items), ctl.slots, ctl.dispatch, next_result, cleanup), 10)
assert self.api.submit.await_count == len(self.state.items)
assert [row["text"] for row in results] == [str(i) for i in range(len(self.state.items))]
assert peak == 16
assert not self.server.distributed_pending_jobs
assert all(task.done() for task in self.manager.tasks.values())
async def test_waits_for_task_drain_not_provider_terminal_or_queue_message(self):
entered, release = asyncio.Event(), asyncio.Event()
original = self.manager._notify_capacity
async def delayed():
entered.set()
await release.wait()
await original()
with patch.object(self.manager, "_notify_capacity", delayed):
slot, handle = await self.dispatch()
self.gates["provider-0"].set()
await entered.wait()
assert self.server.distributed_pending_jobs[handle["job_id"]].qsize() == 1
assert self.controller.poll({slot: handle}) is None
with pytest.raises(RuntimeError, match="occupied"):
await self.controller.dispatch(slot, 1)
release.set()
await handle["task"]
assert self.controller.poll({slot: handle})[1] == 0
async def test_unknown_reservations_survive_and_prevent_provider_admission(self):
for i in range(16):
self.manager.ledger.update(f"unknown-{i}", worker_id="pool", status="UNKNOWN", job_id=None)
slot, handle = await self.dispatch()
await asyncio.sleep(0)
self.api.submit.assert_not_awaited()
assert self.controller.poll({slot: handle}) is None
await self.controller.close()
assert sum(r["status"] == "UNKNOWN" for r in self.manager.ledger.all()) == 16
assert not self.server.distributed_pending_jobs
assert self.manager.pool_is_active("pool")
async def test_local_pending_cap_checked_before_journal(self):
waiting = [asyncio.create_task(asyncio.Event().wait()) for _ in range(64)]
self.manager.tasks.update({f"unrelated-{i}": t for i, t in enumerate(waiting)})
self.manager.runs.update({f"unrelated-{i}": {"execution_id": "unrelated"} for i in range(64)})
try:
with pytest.raises(ValueError, match="limit 64"):
await self.dispatch()
assert not self.manager.ledger.all()
self.api.submit.assert_not_awaited()
assert not self.server.distributed_pending_jobs
finally:
for i, task in enumerate(waiting):
self.manager.tasks.pop(f"unrelated-{i}")
self.manager.runs.pop(f"unrelated-{i}")
task.cancel()
await asyncio.gather(*waiting, return_exceptions=True)
async def test_missing_corrupt_or_mismatched_results_are_atomic_errors(self):
mutations = [lambda o: o["media"]["items"].pop(),
lambda o: o["media"]["artifacts"][-1].update(sha256="0" * 64),
lambda o: o["media"].update(request_id="wrong-request"),
lambda o: o["media"]["bindings"][-1].update(participant_id="wrong-worker")]
for index, mutation in enumerate(mutations):
slot, handle = await self.dispatch(index)
self.mutations[f"provider-{index}"] = mutation
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
assert queue.qsize() == 1
with pytest.raises(RuntimeError, match="failed"):
self.controller.poll({slot: handle})
async def test_terminal_queue_identity_completeness_and_multiplicity(self):
for index, kind in enumerate(("missing", "duplicate", "worker", "mask", "text", "multiple", "empty")):
slot, handle = await self.dispatch(index)
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
message = queue.get_nowait()
if kind == "duplicate":
queue.put_nowait(message)
if kind == "worker":
message["worker_id"] = "unrelated"
if kind in ("mask", "text"):
message["media"].pop("masks" if kind == "mask" else "texts")
if kind == "multiple":
message["media"]["images"].append(item()["image"])
if kind == "empty":
message["media"]["texts"] = [None]
if kind != "missing":
queue.put_nowait(message)
with pytest.raises(RuntimeError):
self.controller.poll({slot: handle})
async def test_wrong_frame_count_channels_and_unpaired_mask_fail(self):
for index, kind in enumerate(("batch", "channels", "grid")):
slot, handle = await self.dispatch(index)
await self.finish(handle)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
message = queue.get_nowait()
if kind == "batch":
message["media"]["images"] = [torch.ones(2, 2, 3, 4)]
elif kind == "channels":
message["media"]["images"] = [torch.ones(1, 2, 3, 1)]
else:
message["media"]["masks"] = [torch.ones(1, 5, 5)]
queue.put_nowait(message)
with pytest.raises(ValueError):
self.controller.poll({slot: handle})
async def test_close_cancels_only_owned_execution_and_exact_queues(self):
state = copy.copy(self.state)
state.context_id = "unrelated"
unrelated = await self.controller_for(state=state)
other_slot, other = await self.dispatch(0, unrelated)
slot, handle = await self.dispatch(0)
queue = self.server.distributed_pending_jobs[handle["job_id"]]
replacement = asyncio.Queue()
self.server.distributed_pending_jobs[handle["job_id"]] = replacement
await self.controller.close()
assert other["task"].done() is False
assert self.server.distributed_pending_jobs[handle["job_id"]] is replacement
assert other["job_id"] in self.server.distributed_pending_jobs
assert handle["task"].done()
assert len(self.api.cancel.await_args_list) == 1
assert queue.empty()
self.server.distributed_pending_jobs.pop(handle["job_id"])
# A late manager callback skips a retired destination instead of recreating it.
run = next(r for r in self.manager.runs.values() if r["execution_id"] == handle["execution_id"])
await self.manager._deliver(run, error="late")
assert handle["job_id"] not in self.server.distributed_pending_jobs
async def test_cancel_dispatch_after_start_loses_handle_but_not_ownership(self):
entered = asyncio.Event()
original = self.manager.start
async def delayed(*args):
ids = await original(*args)
entered.set()
await asyncio.Event().wait()
return ids
with patch.object(self.manager, "start", delayed):
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
await entered.wait()
await self.started.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await self.controller.close()
assert all(t.done() for t in self.manager.tasks.values())
assert not self.server.distributed_pending_jobs
self.api.cancel.assert_awaited_once_with("provider-0")
async def test_close_during_preparation_prevents_late_start(self):
entered = asyncio.Event()
prepared = await prepare()
async def resistant(*args, **kwargs):
entered.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
return prepared
with patch.object(adapter, "prepare_serverless_item", resistant):
task = asyncio.create_task(self.controller.dispatch(self.controller.slots[0], 0))
await entered.wait()
await self.controller.close()
with pytest.raises(RuntimeError, match="closed"):
await task
assert not self.manager.ledger.all()
assert not self.server.distributed_pending_jobs
self.api.submit.assert_not_awaited()
async def test_cancel_during_provider_submission_preserves_unknown(self):
entered = asyncio.Event()
async def ambiguous(*args, **kwargs):
entered.set()
await asyncio.Event().wait()
self.api.submit.side_effect = ambiguous
slot, handle = await self.dispatch()
await entered.wait()
await self.controller.close()
assert next(iter(self.manager.ledger.all()))["status"] == "UNKNOWN"
assert self.manager.pool_is_active("pool")
self.api.cancel.assert_not_awaited()
assert not self.server.distributed_pending_jobs
async def test_cancelled_wait_closes_owned_work(self):
entered = asyncio.Event()
async def job():
try:
await self.dispatch()
entered.set()
await asyncio.Event().wait()
finally:
await self.controller.close()
task = asyncio.create_task(job())
await entered.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert all(t.done() for t in self.manager.tasks.values())
assert not self.server.distributed_pending_jobs
async def test_repeated_cancellation_does_not_abandon_close(self):
slot, handle = await self.dispatch()
entered, release = asyncio.Event(), asyncio.Event()
original = self.manager.cancel
async def delayed(**kwargs):
entered.set()
await release.wait()
return await original(**kwargs)
with patch.object(self.manager, "cancel", delayed):
closing = asyncio.create_task(self.controller.close())
await entered.wait()
closing.cancel()
await asyncio.sleep(0)
closing.cancel()
await asyncio.sleep(0)
assert not closing.done()
assert not self.server.distributed_pending_jobs
release.set()
with pytest.raises(asyncio.CancelledError):
await closing
assert handle["task"].done()
assert self.controller._close_task.done()
async def test_serialization_failure_never_journals_or_submits(self):
with patch.object(adapter, "serialize_request", side_effect=ValueError("budget")):
with pytest.raises(ValueError, match="budget"):
await self.dispatch()
assert not self.manager.ledger.all()
assert not self.server.distributed_pending_jobs
self.api.submit.assert_not_awaited()
async def test_duplicate_index_and_mismatched_handle_fail(self):
slot, handle = await self.dispatch()
with pytest.raises(ValueError, match="duplicate"):
await self.controller.dispatch(self.controller.slots[1], 0)
with pytest.raises(RuntimeError, match="Mismatched"):
self.controller.poll({slot: dict(handle)})
await self.finish(handle)
self.controller.poll({slot: handle})
with pytest.raises(RuntimeError, match="duplicate"):
self.controller.poll({slot: handle})
@@ -0,0 +1,359 @@
"""Offline image-list queue/scheduler integration; no provider or GPU calls."""
import asyncio
import copy
import json
import os
from pathlib import Path
import subprocess
import sys
import types
import unittest
from unittest.mock import AsyncMock, patch
import torch
import test_image_list_master_runtime as direct_fixture
from test_serverless_integration import PKG, load_modules
def pool(**updates):
return dict({"id": "pool", "name": "Cloud", "type": "serverless", "enabled": True,
"provider": "runpod", "endpoint_id": "endpoint", "jobs_per_run": 16,
"timeout_seconds": 30, "media_profile": "comfy-media-v1"}, **updates)
class MixedRuntimeTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Reuse only fixture setup, not inherited tests. The transport/slot doubles
# execute real runtime staging, assignment and result ordering code.
self.fixture = direct_fixture.MasterRuntimeTests()
self.fixture.setUp()
# This reused TestCase has no asyncio runner: invoke its patch cleanup
# directly, not doCleanups(), which would silently fail and leak aiohttp.
self.addCleanup(self.fixture.modules.stop)
self.f = f = self.fixture
self.rt = f.rt
self.configured_pool = pool()
f.config["workers"] = [self.configured_pool]
self.cloud_calls = []
self.completed = []
self.closed = 0
self.hold = False
self.missing = False
self.dispatch_error = False
self.poll_entered = asyncio.Event()
self.dispatch_entered = asyncio.Event()
self.hold_dispatch = False
self.slot_count = 2
fixture = self
class Controller:
def __init__(self, state, event):
self.state, self.event = state, event
self.slots = [("serverless", "pool", index) for index in range(fixture.slot_count)]
self.tasks = []
self.closed = False
async def dispatch(self, slot, index):
assert not self.closed
fixture.cloud_calls.append((slot, index))
fixture.dispatch_entered.set()
if fixture.dispatch_error:
raise RuntimeError("serverless dispatch failed")
async def execute():
if fixture.hold:
await asyncio.Event().wait()
# Make a later source item complete before the first local one.
await asyncio.sleep(0.001)
fixture.completed.append(index)
if fixture.missing:
raise RuntimeError("serverless completed without an image result")
item = state.items[index]
return {"image": item["image"], "mask": item["mask"],
"text": item["values"]["value_1"]}
state = self.state
task = asyncio.create_task(execute())
self.tasks.append(task)
task.add_done_callback(lambda _: self.event.set())
# Own the task before any await that might be cancelled.
if fixture.hold_dispatch:
await asyncio.Event().wait()
return {"index": index, "task": task}
def poll(self, inflight):
fixture.poll_entered.set()
for slot, handle in inflight.items():
if slot in self.slots and handle["task"].done():
return slot, handle["index"], handle["task"].result()
return None
async def close(self):
self.closed = True
for task in self.tasks:
if not task.done():
task.cancel()
await asyncio.gather(*self.tasks, return_exceptions=True)
fixture.closed += 1
async def create(state, region, config, event):
self.assertEqual(region["end_id"], "end")
self.assertEqual(config["workers"][-1]["id"], "pool")
self.controller = Controller(state, event)
return self.controller
self.create_controller = AsyncMock(side_effect=create)
sys.modules['_master_slot_tests.image_list.serverless'] = types.SimpleNamespace(
create_serverless_slots=self.create_controller)
self.direct_dispatch = sys.modules['_master_slot_tests.api.orchestration.dispatch'].dispatch_worker_prompt
def create(self, *, master=True, direct=True, count=9):
f = self.f
if direct:
f.workers.append({"id": "peer", "type": "local"})
f.config["workers"] = f.workers + [self.configured_pool]
f.local_delay = 0.04
context = f.create(participates=master, count=count) if (master or direct) else None
if context is None:
# Fixture.create assumes only direct peers; cloud-only uses a pool ID.
f.items = self.rt.normalize_image_list_inputs(torch.zeros((count, 2, 2, 3)),
mask=torch.zeros((count, 2, 2)), values={"value_1": [str(i) for i in range(count)]})
context = self.rt.create_master_job('start', 'end', direct_fixture.graph(), f.items, ['pool'], '')
f.context = context
f.state = self.rt.get_master_job(context)
else:
f.state.worker_ids.append('pool')
return context
async def test_mixed_roles_share_one_queue_and_preserve_pairing(self):
context = self.create()
original = copy.deepcopy(self.f.state.original_prompt)
result = await self.rt.run_image_list_job(context)
indices = self.f.local_calls + self.f.dispatched + [i for _, i in self.cloud_calls]
self.assertEqual(sorted(indices), list(range(9)))
self.assertTrue(self.f.local_calls and self.f.dispatched and self.cloud_calls)
self.assertNotEqual(self.completed[0], 0)
self.assertEqual(result[2], [str(i) for i in range(9)])
for index, (image, mask) in enumerate(zip(*result[:2])):
self.assertTrue(torch.equal(image, self.f.items[index]['image']))
self.assertTrue(torch.equal(mask, self.f.items[index]['mask']))
self.assertEqual(self.f.state.original_prompt, original)
self.f.preflight.assert_awaited_once_with(self.f.workers, require_idle=False)
for call in self.direct_dispatch.await_args_list:
self.assertEqual(call.kwargs['partial_execution_targets'], ['end'])
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_cloud_only_one_job_per_item_ignores_images_per_run(self):
for jobs_per_run in (1, 16):
with self.subTest(jobs_per_run=jobs_per_run):
self.configured_pool['jobs_per_run'] = jobs_per_run
self.cloud_calls.clear()
result = await self.rt.run_image_list_job(self.create(master=False, direct=False, count=7))
self.assertEqual([i for _, i in self.cloud_calls], list(range(7)))
self.assertEqual(result[2], [str(i) for i in range(7)])
self.assertEqual(len({slot for slot, _ in self.cloud_calls}), self.slot_count)
self.f.preflight.assert_not_awaited()
self.f.network.get_client_session.assert_not_awaited()
self.assertFalse(self.f.slot_instances)
async def test_missing_cloud_result_fails_without_refill(self):
self.missing = True
with self.assertRaisesRegex(RuntimeError, 'without an image result'):
await self.rt.run_image_list_job(self.create(master=False, direct=False))
self.assertEqual(len(self.cloud_calls), self.slot_count)
self.assertEqual(self.closed, 1)
self.assertFalse(self.rt._MASTER_JOBS)
async def test_controller_closed_when_direct_preflight_fails(self):
self.f.preflight.side_effect = RuntimeError('direct preflight failed')
with self.assertRaisesRegex(RuntimeError, 'direct preflight failed'):
await self.rt.run_image_list_job(self.create())
self.assertEqual(self.closed, 1)
self.assertFalse(self.cloud_calls)
async def test_controller_closed_when_local_setup_fails(self):
slot_module = sys.modules['_master_slot_tests.image_list.local_executor']
with patch.object(slot_module, 'LocalExecutionSlot', side_effect=RuntimeError('local setup failed')):
with self.assertRaisesRegex(RuntimeError, 'local setup failed'):
await self.rt.run_image_list_job(self.create())
self.assertEqual(self.closed, 1)
self.assertFalse(self.cloud_calls)
async def test_controller_closed_on_dispatch_failure(self):
self.dispatch_error = True
with self.assertRaisesRegex(RuntimeError, 'serverless dispatch failed'):
await self.rt.run_image_list_job(self.create(master=False, direct=False))
self.assertEqual(len(self.cloud_calls), 1)
self.assertEqual(self.closed, 1)
async def test_cancelled_dispatch_without_returned_handle_is_owned(self):
self.hold = self.hold_dispatch = True
with patch.object(self.rt, 'SUBMISSION_CANCEL_GRACE_SECONDS', 0.01):
task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False)))
await asyncio.wait_for(self.dispatch_entered.wait(), 1)
task.cancel()
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 1)
self.assertEqual(len(self.cloud_calls), 1)
self.assertEqual(self.closed, 1)
self.assertTrue(all(t.done() for t in self.controller.tasks))
async def interrupt_probe(self, models):
self.hold = True
task = asyncio.create_task(self.rt.run_image_list_job(self.create(master=False, direct=False)))
try:
await asyncio.wait_for(self.poll_entered.wait(), 1)
assigned = list(self.cloud_calls)
models.interrupt_current_processing(True)
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 1)
self.assertEqual(self.cloud_calls, assigned)
self.assertEqual(self.closed, 1)
self.assertTrue(all(t.done() for t in self.controller.tasks))
self.assertFalse(self.rt._MASTER_JOBS)
finally:
models.interrupt_current_processing(False)
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
class QueueSelectionTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.q, modules, self.ps = load_modules()
self.modules_patch = patch.dict(sys.modules, modules)
self.modules_patch.start()
self.addCleanup(self.modules_patch.stop)
self.pool = pool()
self.peer = {'id': 'peer', 'name': 'Peer', 'type': 'local', 'enabled': True}
self.config = {'workers': [self.pool], 'settings': {'master_delegate_only': True}}
self.q.load_config = lambda: self.config
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw: (workers, delegate))
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
sys.modules[PKG + '.api.orchestration.dispatch'].preflight_image_list_workers = self.preflight
self.q.queue_prompt_payload = AsyncMock(return_value={'prompt_id': 'outer', 'number': 1})
self.prompt = direct_fixture.graph()
async def queue(self, **kwargs):
return await self.q.orchestrate_distributed_execution(self.prompt, None, 'client', **kwargs)
async def test_delegate_pool_selected_once_no_direct_probe_or_source_mutation(self):
original = copy.deepcopy(self.prompt)
result = await self.queue()
selected = self.q.select_active_workers.await_args.args[0]
self.assertEqual([w['id'] for w in selected], ['pool'])
prepared = self.q.queue_prompt_payload.await_args.args[0]
inputs = prepared['start']['inputs']
self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['pool'])
self.assertIs(inputs['image_list_master_participates'], False)
self.assertEqual(self.q.queue_prompt_payload.await_args.kwargs['partial_execution_targets'], ['end'])
self.assertEqual(result[2], 1)
self.assertEqual(self.prompt, original)
self.preflight.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
async def test_mixed_selection_preflights_only_direct_peers(self):
self.config['workers'].insert(0, self.peer)
result = await self.queue(delegate_master=False)
self.assertEqual([w['id'] for w in self.preflight.await_args.args[0]], ['peer'])
self.assertFalse(self.preflight.await_args.kwargs['require_idle'])
inputs = self.q.queue_prompt_payload.await_args.args[0]['start']['inputs']
self.assertEqual(json.loads(inputs['image_list_enabled_worker_ids']), ['peer', 'pool'])
self.assertEqual(result[2], 2)
async def test_legacy_pool_rejected_before_probe_or_queue(self):
self.pool['media_profile'] = 'legacy-image-v1'
with self.assertRaisesRegex((ValueError, RuntimeError), 'comfy-media-v1'):
await self.queue()
self.q.select_active_workers.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_unknown_selection_rejected_before_queue(self):
with self.assertRaisesRegex((ValueError, RuntimeError), 'Unknown|unknown'):
await self.queue(enabled_worker_ids=['missing'])
self.q.select_active_workers.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_disabled_pool_can_be_explicitly_selected(self):
self.pool['enabled'] = False
await self.queue(enabled_worker_ids=['pool'])
self.assertEqual([w['id'] for w in self.q.select_active_workers.await_args.args[0]], ['pool'])
async def test_explicit_empty_selection_does_not_import_cloud_adapter(self):
sys.modules[PKG + '.image_list.serverless'] = None
result = await self.queue(enabled_worker_ids=[], delegate_master=False)
self.assertEqual(result[2], 0)
self.preflight.assert_not_awaited()
async def test_cloud_auth_or_capacity_failure_never_queues_master(self):
for error in ('authentication failed', 'capacity exhausted'):
with self.subTest(error=error):
self.q.select_active_workers.side_effect = ValueError(error)
with self.assertRaisesRegex(ValueError, error):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_real_selection_uses_manager_health_not_direct_http(self):
from test_dispatch_selection import _load_dispatch_module
dispatch = _load_dispatch_module()
health = AsyncMock(return_value={'ok': True, 'queueRemaining': 3})
sys.modules['dist_dispatch_testpkg.api.orchestration.serverless_runtime'] = types.SimpleNamespace(
get_manager=lambda: types.SimpleNamespace(health=health))
self.q.select_active_workers = dispatch.select_active_workers
with patch.object(dispatch, 'get_client_session', AsyncMock(side_effect=AssertionError('no HTTP'))), \
patch.object(dispatch, 'probe_worker', AsyncMock(side_effect=AssertionError('no direct probe'))):
await self.queue()
health.assert_awaited_once()
self.assertEqual(health.await_args.args[0]['id'], 'pool')
self.preflight.assert_not_awaited()
self.q.queue_prompt_payload.reset_mock()
health.side_effect = ValueError('provider authentication failed')
with self.assertRaisesRegex(ValueError, 'authentication failed'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_unreachable_pool_never_enables_master_fallback(self):
self.q.select_active_workers.side_effect = None
self.q.select_active_workers.return_value = ([], False)
with self.assertRaisesRegex(RuntimeError, 'orchestrator|Orchestrator'):
await self.queue()
self.q.queue_prompt_payload.assert_not_awaited()
class RealInterruptTests(unittest.TestCase):
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for CPU interrupt probe')
def test_processing_interrupted_signal_during_cloud_result_wait(self):
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--interrupt-probe',
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=45)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn('SERVERLESS_INTERRUPT_OK', result.stdout)
def real_interrupt_probe(comfy_root):
sys.path.insert(0, str(Path(comfy_root).resolve()))
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import comfy.model_management as models
real_comfy = sys.modules['comfy']
case = MixedRuntimeTests()
case.setUp()
try:
# Keep actual thread-safe ComfyUI interruption state, not a task.cancel
# surrogate or a fake processing_interrupted callback.
sys.modules['comfy'] = real_comfy
sys.modules['comfy.model_management'] = models
asyncio.run(case.interrupt_probe(models))
finally:
case.f.modules.stop()
print('SERVERLESS_INTERRUPT_OK no refill, owned close, actual ComfyUI flag')
if __name__ == '__main__':
if len(sys.argv) > 1 and sys.argv[1] == '--interrupt-probe':
real_interrupt_probe(sys.argv[2])
else:
unittest.main()
+106
View File
@@ -0,0 +1,106 @@
import importlib.util
import sys
import types
from pathlib import Path
import pytest
def _load_config_module():
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
package_name = "image_list_timeout_config_testpkg"
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__ = [str(module_path.parent)]
sys.modules[package_name] = root_pkg
logging_module = types.ModuleType(f"{package_name}.logging")
logging_module.log = lambda *_args, **_kwargs: None
logging_module.debug_log = lambda *_args, **_kwargs: None
sys.modules[f"{package_name}.logging"] = logging_module
constants_module = types.ModuleType(f"{package_name}.constants")
constants_module.HEARTBEAT_TIMEOUT = 30
sys.modules[f"{package_name}.constants"] = constants_module
spec = importlib.util.spec_from_file_location(f"{package_name}.config", 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)
return module
config_module = _load_config_module()
SETTING_KEY = "image_list_job_timeout_seconds"
def _config_with_timeout(value):
return {"settings": {SETTING_KEY: value}}
def test_default_config_sets_image_list_job_timeout_to_zero():
assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0
def test_missing_value_defaults_to_disabled_timeout():
assert config_module.get_image_list_job_timeout_seconds({}) == 0.0
assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0
def test_none_config_loads_current_config(monkeypatch):
cfg = _config_with_timeout("12.5")
monkeypatch.setattr(config_module, "load_config", lambda: cfg)
assert config_module.get_image_list_job_timeout_seconds() == 12.5
@pytest.mark.parametrize(
("value", "expected"),
[
(0, 0.0),
(0.0, 0.0),
("0", 0.0),
(" 0.0 ", 0.0),
(1, 1.0),
(2.5, 2.5),
("3", 3.0),
(" 4.25 ", 4.25),
],
)
def test_accepts_zero_and_positive_numeric_values(value, expected):
result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
assert result == expected
assert isinstance(result, float)
@pytest.mark.parametrize(
"value",
[
True,
False,
-1,
-0.01,
float("nan"),
float("inf"),
float("-inf"),
"not-a-number",
"",
" ",
"-1",
"nan",
"inf",
None,
[],
{},
],
)
def test_rejects_invalid_values_with_setting_name(value):
with pytest.raises(ValueError, match=SETTING_KEY):
config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
+70
View File
@@ -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()
+265
View File
@@ -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())
+212
View File
@@ -0,0 +1,212 @@
import asyncio
import importlib.util
import sys
import types
import unittest
import uuid
from pathlib import Path
from typing import Any, cast
from unittest.mock import patch
import aiohttp
from aiohttp import web
def _load_dispatch_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}"
package_modules = {}
for suffix in ("", ".api", ".api.orchestration", ".utils"):
package = types.ModuleType(package_name + suffix)
package.__path__ = []
package_modules[package_name + suffix] = package
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
setattr(logging_module, "log", lambda *_args, **_kwargs: None)
package_modules[logging_module.__name__] = logging_module
network_module = types.ModuleType(f"{package_name}.utils.network")
setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid")
async def _unused_probe_worker(*_args, **_kwargs):
return None
async def _unused_get_client_session(*_args, **_kwargs):
raise RuntimeError("get_client_session must be patched by this test")
setattr(network_module, "probe_worker", _unused_probe_worker)
setattr(network_module, "get_client_session", _unused_get_client_session)
package_modules[network_module.__name__] = network_module
module_name = f"{package_name}.api.orchestration.dispatch"
spec = importlib.util.spec_from_file_location(module_name, module_path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
package_modules[module_name] = module
previous_modules = {name: sys.modules.get(name) for name in package_modules}
try:
sys.modules.update(package_modules)
spec.loader.exec_module(module)
finally:
for name, previous in previous_modules.items():
if previous is None:
sys.modules.pop(name, None)
else:
sys.modules[name] = previous
return module
dispatch = _load_dispatch_module()
class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.mode = None
self.received_payload = None
self.dispatch_received = asyncio.Event()
self.peer_finished = asyncio.Event()
self.irrelevant_messages_sent = 0
app = web.Application()
app.router.add_get("/distributed/worker_ws", self._handle_websocket)
self.runner = web.AppRunner(app)
await self.runner.setup()
self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
await self.site.start()
server = cast(Any, self.site._server)
port = server.sockets[0].getsockname()[1]
self.worker_url = f"http://127.0.0.1:{port}"
self.session = aiohttp.ClientSession()
async def asyncTearDown(self):
await self.session.close()
await self.runner.cleanup()
async def _get_client_session(self):
return self.session
async def _handle_websocket(self, request):
ws = web.WebSocketResponse()
await ws.prepare(request)
try:
payload = await ws.receive_json()
self.received_payload = payload
self.dispatch_received.set()
if self.mode in ("success", "slow-close"):
await ws.send_json({
"type": "dispatch_ack",
"request_id": "unrelated-request",
"ok": True,
"prompt_id": "wrong-prompt",
})
await ws.send_json({"type": "heartbeat"})
await ws.send_json({
"type": "dispatch_ack",
"request_id": payload["request_id"],
"ok": True,
"prompt_id": "prompt-123",
})
if self.mode == "slow-close":
await asyncio.sleep(0.15)
elif self.mode == "rejection":
await ws.send_json({
"type": "dispatch_ack",
"request_id": payload["request_id"],
"ok": False,
"error": "Invalid prompt",
"validation_error": {"node": "bad"},
"node_errors": {"bad": "missing input"},
})
elif self.mode in ("no-ack", "wait-for-cancel"):
async def send_irrelevant_messages():
while True:
try:
await ws.send_json({"type": "heartbeat"})
self.irrelevant_messages_sent += 1
await ws.send_json({
"type": "dispatch_ack",
"request_id": "unrelated-request",
"ok": True,
"prompt_id": "not-this-request",
})
self.irrelevant_messages_sent += 1
except (ConnectionError, RuntimeError):
return
await asyncio.sleep(0.01)
sender = asyncio.create_task(send_irrelevant_messages())
try:
async for _message in ws:
pass
finally:
sender.cancel()
try:
await sender
except asyncio.CancelledError:
pass
finally:
self.peer_finished.set()
return ws
async def _dispatch(self, timeout):
with patch.object(dispatch, "get_client_session", self._get_client_session):
return await dispatch._dispatch_via_websocket(
self.worker_url,
{"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}},
"client-456",
timeout=timeout,
)
async def test_returns_matching_acknowledgement_prompt_id(self):
self.mode = "success"
prompt_id = await self._dispatch(timeout=0.5)
self.assertEqual(prompt_id, "prompt-123")
received_payload = self.received_payload
assert received_payload is not None
self.assertEqual(received_payload["type"], "dispatch_prompt")
self.assertTrue(received_payload["request_id"])
self.assertEqual(received_payload["client_id"], "client-456")
self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}})
self.assertEqual(received_payload["workflow"], {"name": "test"})
async def test_acknowledged_prompt_id_survives_slow_close(self):
self.mode = "slow-close"
self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123")
async def test_matching_rejection_preserves_worker_error_details(self):
self.mode = "rejection"
with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="):
await self._dispatch(timeout=0.5)
async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self):
self.mode = "no-ack"
started = asyncio.get_running_loop().time()
with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"):
await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5)
elapsed = asyncio.get_running_loop().time() - started
self.assertLess(elapsed, 0.5)
self.assertGreaterEqual(self.irrelevant_messages_sent, 4)
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
async def test_outer_cancellation_propagates_and_closes_peer(self):
self.mode = "wait-for-cancel"
task = asyncio.create_task(self._dispatch(timeout=2.0))
await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5)
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=0.5)
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
if __name__ == "__main__":
unittest.main()
+198
View File
@@ -0,0 +1,198 @@
"""Real ComfyUI V3 bridge mapping, file ownership and exact CPU round trips."""
import os
from pathlib import Path
import subprocess
import sys
import pytest
ROOT = Path(__file__).resolve().parents[1]
def test_real_media_bridge_nodes():
core = os.environ.get("COMFYUI_SOURCE_ROOT")
if core is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU media bridge acceptance")
assert core and (Path(core) / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
"Configured COMFYUI_SOURCE_ROOT must provide native video classes"
)
probe = r'''
import asyncio, importlib.util, io, json, os, sys, tempfile
from fractions import Fraction
import av
import numpy as np
from pathlib import Path
from unittest.mock import patch
import torch
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
def forbidden(*args, **kwargs):
raise AssertionError('Bridge tests must not query or initialize a GPU')
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
for name in ('init', '_lazy_init', 'current_device'):
setattr(torch.cuda, name, forbidden)
import execution
import folder_paths
from comfy_api.latest._input_impl.video_types import VideoFromComponents, VideoFromFile, VideoFromList
from comfy_api.latest._util import VideoComponents
for name in ('is_available', 'device_count'):
setattr(torch.cuda, name, forbidden)
spec = importlib.util.spec_from_file_location('typed_bridge_probe', Path(sys.argv[2])/'serverless_bridge/__init__.py', submodule_search_locations=[str(Path(sys.argv[2])/'serverless_bridge')])
bridge = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = bridge
spec.loader.exec_module(bridge)
from typed_bridge_probe.nodes import MediaImportV1, MediaExportV1, export_values, import_value
from typed_bridge_probe.contract import SOCKET_KINDS
from typed_bridge_probe.video_codec import encode_video
def tiny_video():
buffer = io.BytesIO()
with av.open(buffer, 'w', format='matroska') as container:
stream = container.add_stream('ffv1', rate=Fraction(24, 1))
stream.width = stream.height = 32
stream.pix_fmt = 'yuv420p'
for level in (16, 64, 128):
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), level, dtype=np.uint8), format='rgb24')
for packet in stream.encode(frame): container.mux(packet)
for packet in stream.encode(): container.mux(packet)
return buffer.getvalue()
def exact_tensor(actual, expected):
assert actual.device.type == 'cpu' and actual.dtype == expected.dtype
assert actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def exact_video(actual, expected, wire):
assert type(actual) is type(expected)
assert encode_video(actual) == wire
if type(actual) is VideoFromComponents:
a = vars(actual)['_VideoFromComponents__components']
e = vars(expected)['_VideoFromComponents__components']
assert type(a) is VideoComponents and type(a.frame_rate) is Fraction
assert a.frame_rate == e.frame_rate and a.metadata == e.metadata
assert a.metadata is not e.metadata
for name in ('images', 'alpha'): exact_tensor(getattr(a, name), getattr(e, name))
exact_tensor(a.audio['waveform'], e.audio['waveform'])
assert a.audio['sample_rate'] == e.audio['sample_rate']
assert actual.get_bit_depth() == 10 and actual.get_color_space() == 'HDR PQ'
else:
values = actual.videos if type(actual) is VideoFromList else [actual]
for value in values:
buffer = vars(value)['_VideoFromFile__file']
assert type(buffer) is io.BytesIO and not buffer.closed
buffer.seek(0)
with av.open(buffer) as container:
assert len(list(container.decode(video=0))) == 3
if type(actual) is VideoFromList:
exact_tensor(actual.complete_audio['waveform'], expected.complete_audio['waveform'])
async def check():
extension = await bridge.comfy_entrypoint()
assert await extension.get_node_list() == [MediaImportV1, MediaExportV1]
imp = MediaImportV1.GET_NODE_INFO_V1()
exp = MediaExportV1.GET_NODE_INFO_V1()
assert imp['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], imp
assert exp['output_node'] and MediaExportV1.define_schema().is_input_list, exp
with tempfile.TemporaryDirectory(prefix='typed-media-', dir=os.environ.get('TMPDIR')) as scratch:
base = Path(scratch)
roots = {key: base/key for key in ('input','output','temp')}
for path in roots.values(): path.mkdir()
def annotated(name):
for key, root in roots.items():
suffix = f' [{key}]'
if name.endswith(suffix): return str(root/name[:-len(suffix)])
return str(roots['input']/name)
with patch.object(folder_paths,'get_input_directory',lambda: str(roots['input'])), patch.object(folder_paths,'get_output_directory',lambda: str(roots['output'])), patch.object(folder_paths,'get_temp_directory',lambda: str(roots['temp'])), patch.object(folder_paths,'get_annotated_filepath',annotated):
examples = {'images': torch.rand(2,2,3,4,dtype=torch.float64),
'masks': torch.tensor([[[0.123456789,0.987654321]]],dtype=torch.float64),
'audio': {'waveform':torch.randn(1,6,31,dtype=torch.float64),'sample_rate':44100},
'texts': ' café\n🎥\r\n'}
media = tiny_video()
sources = [roots['input']/f'source-{n}.mkv' for n in range(3)]
for path in sources: path.write_bytes(media)
images = torch.arange(36, dtype=torch.float64, device='cpu').reshape(2,2,3,3) / 997
images[0,0,0,0] = -0.0
component_video = VideoFromComponents(VideoComponents(
images=images, frame_rate=Fraction(30000,1001),
alpha=torch.full((2,2,3,1),0.123456789,dtype=torch.float64,device='cpu'),
audio=examples['audio'], metadata={'title':'雪', 'nested':[True,None,-4,0.125]}),
bit_depth=10,color_space='HDR PQ')
videos = [component_video,
VideoFromFile(str(sources[0]),start_time=-0.0,duration=0,crop=(0,0,2,2)),
VideoFromList([VideoFromFile(str(sources[1]),start_time=-1.25,duration=-2.5,crop=(-3,5,7,-9)),
VideoFromFile(str(sources[2]),start_time=0,duration=0)],
complete_audio=examples['audio'])]
for socket,value in [*examples.items(), *[('video',v) for v in videos]]:
kind = SOCKET_KINDS[socket]
expected_wire = encode_video(value) if kind == 'VIDEO' else None
binding = dict(binding_id='b-'+socket,collector_id='collector',job_id='job',participant_id='worker',node_id='export-'+socket,socket=socket,kind=kind)
controls = {'request_id':'request','binding_json':json.dumps(binding)}
data, missing, v3 = execution.get_input_data(controls,MediaExportV1,'export-'+socket)
assert not missing, missing
data[socket] = [value]
results = await execution._async_map_node_over_list('probe','export-'+socket,MediaExportV1,data,MediaExportV1.FUNCTION,v3_data=v3)
results = await execution.resolve_map_node_over_list_results(results)
raw = results[0].ui
exported = raw['media_exports'][0]
assert exported['binding']['count'] == 1
artifact = exported['artifacts'][0]
transport = artifact.pop('transport')
filename = transport['subfolder']+'/'+transport['filename']+' [temp]'
decoded = import_value(artifact,filename)
if kind in ('IMAGE','MASK'):
assert decoded.dtype == value.dtype and torch.equal(decoded,value)
elif kind == 'AUDIO':
assert decoded['sample_rate'] == value['sample_rate']
assert torch.equal(decoded['waveform'],value['waveform'])
elif kind == 'VIDEO': exact_video(decoded,value,expected_wire)
else: assert decoded == value
mapped, missing, hidden = execution.get_input_data({'artifact_json':json.dumps(artifact),'filename':filename},MediaImportV1,'import')
got = await execution._async_map_node_over_list('probe','import',MediaImportV1,mapped,MediaImportV1.FUNCTION,v3_data=hidden)
got = await execution.resolve_map_node_over_list_results(got)
assert len(got[0].result) == 5
assert all(x is None for n,x in enumerate(got[0].result) if n != list(SOCKET_KINDS).index(socket))
bad = dict(artifact,sha256='0'*64)
try: import_value(bad,filename)
except ValueError: pass
else: raise AssertionError('checksum accepted')
try: import_value(artifact,str(base/'outside.bin'))
except ValueError: pass
else: raise AssertionError('outside path accepted')
if kind == 'VIDEO':
mapped_video = got[0].result[list(SOCKET_KINDS).index(socket)]
# Remove both the exported artifact and the original input
# containers before checking independently owned BytesIOs.
Path(annotated(filename)).unlink()
originals = value.videos if type(value) is VideoFromList else [value]
for original in originals:
if type(original) is VideoFromFile:
Path(vars(original)['_VideoFromFile__file']).unlink()
exact_video(decoded,value,expected_wire)
exact_video(mapped_video,value,expected_wire)
if type(value) is VideoFromFile:
assert tuple(decoded.get_components().images.shape) == (3,2,2,3)
text_binding = dict(binding_id='empty',collector_id='collector',job_id='job',participant_id='worker',node_id='export-empty',socket='texts',kind='STRING')
out = export_values('request',text_binding,[''])
assert out['artifacts'][0]['byte_length'] == 0
before = set(roots['temp'].rglob('*.bin'))
try: export_values('request',text_binding,['first',None])
except ValueError: pass
else: raise AssertionError('invalid text accepted')
assert set(roots['temp'].rglob('*.bin')) == before
print('MEDIA_BRIDGE_CPU_OK')
asyncio.run(check())
'''
env = dict(os.environ)
env.pop("PYTHONPATH", None)
env.pop("COMFYUI_SOURCE_ROOT", None)
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", PYTHONDONTWRITEBYTECODE="1",
TMPDIR=str(scratch))
result = subprocess.run([sys.executable, "-c", probe, str(Path(core).resolve()), str(ROOT)],
capture_output=True, text=True, timeout=90, env=env)
assert result.returncode == 0, result.stdout + result.stderr
assert "MEDIA_BRIDGE_CPU_OK" in result.stdout
+328
View File
@@ -0,0 +1,328 @@
import importlib.util
import math
import struct
import types
from pathlib import Path
import pytest
import torch
MODULE_PATH = Path(__file__).resolve().parents[1] / "serverless_bridge" / "codec.py"
spec = importlib.util.spec_from_file_location("dist_media_codec_test", MODULE_PATH)
codec = importlib.util.module_from_spec(spec)
spec.loader.exec_module(codec)
DTYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
SHAPES = [
("IMAGE", (2, 3, 5, 1), "BHWC"),
("IMAGE", (2, 3, 5, 3), "BHWC"),
("IMAGE", (2, 3, 5, 4), "BHWC"),
("MASK", (2, 3, 5), "BHW"),
("AUDIO", (2, 1, 17), "BCT"),
("AUDIO", (2, 6, 17), "BCT"),
]
def audio(tensor, sample_rate=48000):
return {"waveform": tensor, "sample_rate": sample_rate}
def descriptor(_kind="MASK", **changes):
result = {
"encoding": "tensor-le-v1", "dtype": "float32",
"shape": [1, 1, 1], "layout": "BHW",
}
if _kind == "IMAGE":
result.update(shape=[1, 1, 1, 3], layout="BHWC")
elif _kind == "AUDIO":
result.update(layout="BCT", sample_rate=48000)
result.update(changes)
return result
@pytest.fixture(autouse=True)
def forbid_gpu_queries(monkeypatch):
def forbidden(*args, **kwargs):
pytest.fail("codec must not query or initialize a GPU")
for name in ("is_available", "init", "_lazy_init", "device_count"):
monkeypatch.setattr(torch.cuda, name, forbidden)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kind,shape,layout", SHAPES)
def test_exact_roundtrip_preserves_dtype_shape_channels_and_bits(dtype, kind, shape, layout):
count = math.prod(shape)
tensor = (torch.arange(count, device="cpu", dtype=torch.float64) / 997 - 0.125).to(dtype).reshape(shape)
tensor.flatten()[0] = -0.0
tensor.requires_grad_(True)
value = audio(tensor) if kind == "AUDIO" else tensor
metadata, payload = codec.encode_payload(kind, value)
expected = {
"encoding": "tensor-le-v1", "dtype": str(dtype).removeprefix("torch."),
"shape": list(shape), "layout": layout,
}
if kind == "AUDIO":
expected["sample_rate"] = 48000
assert metadata == expected
assert type(payload) is bytes
assert len(payload) == count * tensor.element_size()
result = codec.decode_payload(kind, metadata, payload)
if kind == "AUDIO":
assert set(result) == {"waveform", "sample_rate"}
assert result["sample_rate"] == 48000
result = result["waveform"]
assert result.device.type == "cpu"
assert result.is_contiguous()
assert not result.requires_grad
assert result.dtype == dtype
assert result.shape == shape
assert torch.equal(result.view(torch.uint8), tensor.detach().view(torch.uint8))
@pytest.mark.parametrize("dtype,fmt,expected_bits", [
(torch.float16, "e", None), (torch.bfloat16, "H", [0x3F80, 0x8000, 0xC020]),
(torch.float32, "f", None), (torch.float64, "d", None),
])
def test_wire_bytes_are_little_endian(dtype, fmt, expected_bits):
values = [1.0, -0.0, -2.5]
tensor = torch.tensor(values, dtype=dtype, device="cpu").reshape(1, 1, 3)
metadata, payload = codec.encode_payload("MASK", tensor)
expected = struct.pack("<3" + fmt, *(expected_bits if expected_bits else values))
assert payload == expected
assert torch.equal(codec.decode_payload("MASK", metadata, expected).view(torch.uint8), tensor.view(torch.uint8))
def test_bfloat16_extremes_and_subnormal_bits_are_lossless():
bits = torch.tensor([0, 0x8000, 1, 0x8001, 0x3F81, 0x7F7F, 0xFF7F], dtype=torch.uint16, device="cpu")
value = bits.view(torch.bfloat16).reshape(1, 1, -1)
metadata, payload = codec.encode_payload("MASK", value)
assert payload == struct.pack("<7H", *bits.tolist())
assert torch.equal(codec.decode_payload("MASK", metadata, payload).view(torch.uint16).flatten(), bits)
def test_fractional_mask_is_not_quantized_clipped_or_normalized():
value = torch.tensor([0.1234567890123, 1.0000000000001, -2.5], dtype=torch.float64, device="cpu").reshape(1, 1, 3)
metadata, payload = codec.encode_payload("MASK", value)
assert torch.equal(codec.decode_payload("MASK", metadata, payload), value)
@pytest.mark.parametrize("text", ["", " ", "\t\n ", "café 雪 🎵", "e\u0301", "\x00\ufeff\r\n"])
def test_strings_preserve_exact_utf8(text):
metadata, payload = codec.encode_payload("STRING", text)
assert metadata == {"encoding": "utf8-v1"}
assert payload == text.encode("utf-8")
assert codec.decode_payload("STRING", metadata, payload) == text
def test_noncontiguous_and_lazy_negative_tensors_roundtrip():
value = torch.arange(24, dtype=torch.float32, device="cpu").reshape(2, 3, 4).transpose(1, 2)
for tensor in (value, torch._neg_view(value)):
metadata, payload = codec.encode_payload("MASK", tensor)
result = codec.decode_payload("MASK", metadata, payload)
assert result.is_contiguous()
assert torch.equal(result, tensor)
def test_decode_is_independent_and_ignores_default_device(monkeypatch):
source = torch.tensor([0.25], device="cpu").reshape(1, 1, 1)
metadata, payload = codec.encode_payload("MASK", source)
buffers = []
original = torch.frombuffer
def capture(buffer, *args, **kwargs):
buffers.append(buffer)
return original(buffer, *args, **kwargs)
monkeypatch.setattr(torch, "frombuffer", capture)
with torch.device("meta"):
first = codec.decode_payload("MASK", metadata, payload)
second = codec.decode_payload("MASK", metadata, payload)
assert first.device.type == second.device.type == "cpu"
first.fill_(9)
source.fill_(7)
for buffer in buffers:
if isinstance(buffer, bytearray):
buffer[:] = b"\x00" * len(buffer)
assert second.item() == 0.25
assert codec.decode_payload("MASK", metadata, payload).item() == 0.25
@pytest.mark.parametrize("kind", ["image", " MASK", "LATENT", "", None, 1, [], {}])
def test_unknown_kinds_raise_value_error(kind):
with pytest.raises(ValueError, match="kind"):
codec.encode_payload(kind, "")
with pytest.raises(ValueError, match="kind"):
codec.decode_payload(kind, {}, b"")
@pytest.mark.parametrize("kind,value", [
("STRING", None), ("STRING", b"text"), ("STRING", 123),
("IMAGE", [[[[0.5]]]]), ("MASK", None),
("AUDIO", torch.ones(1, 1, 1, device="cpu")), ("AUDIO", {}),
("AUDIO", {"waveform": "invalid", "sample_rate": 48000}),
("AUDIO", {"waveform": torch.ones(1, 1, 1, device="cpu")}),
])
def test_unsupported_values_raise_value_error(kind, value):
with pytest.raises(ValueError):
codec.encode_payload(kind, value)
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int64, torch.bool, torch.complex64])
def test_unsupported_tensor_dtype(dtype):
with pytest.raises(ValueError, match="dtype"):
codec.encode_payload("MASK", torch.ones(1, 1, 1, dtype=dtype, device="cpu"))
@pytest.mark.parametrize("kind,shape", [
("IMAGE", (1, 2, 3)), ("IMAGE", (1, 1, 1, 2)), ("IMAGE", (1, 1, 1, 5)),
("MASK", (1, 2)), ("MASK", (1, 2, 3, 1)), ("AUDIO", (2, 4)),
("IMAGE", (1, 0, 2, 3)), ("MASK", (0, 1, 1)), ("AUDIO", (1, 0, 4)),
("AUDIO", (1, 2, 0)),
])
def test_invalid_tensor_shape(kind, shape):
tensor = torch.empty(shape, dtype=torch.float32, device="cpu")
with pytest.raises(ValueError, match="shape|rank|channel|dimension"):
codec.encode_payload(kind, audio(tensor) if kind == "AUDIO" else tensor)
def test_rejects_non_cpu_and_sparse_without_transferring(monkeypatch):
meta = torch.empty((1, 1, 1), device="meta")
sparse = torch.ones((1, 1, 1), device="cpu").to_sparse()
def forbidden(*args, **kwargs):
pytest.fail("codec must not move tensors between devices")
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
monkeypatch.setattr(torch.Tensor, "to", forbidden)
for tensor in (meta, sparse):
with pytest.raises(ValueError, match="CPU|layout|strided"):
codec.encode_payload("MASK", tensor)
@pytest.mark.parametrize("sample_rate", [True, False, 0, -1, 44100.0, "44100", None])
def test_invalid_sample_rate(sample_rate):
with pytest.raises(ValueError, match="sample_rate"):
codec.encode_payload("AUDIO", audio(torch.ones(1, 1, 1, device="cpu"), sample_rate))
with pytest.raises(ValueError, match="sample_rate"):
codec.decode_payload("AUDIO", descriptor("AUDIO", sample_rate=sample_rate), b"\x00" * 4)
@pytest.mark.parametrize("bad", [None, [], "descriptor", {}, {"encoding": "utf8-v1"}])
def test_malformed_tensor_descriptor(bad):
with pytest.raises(ValueError, match="descriptor"):
codec.decode_payload("MASK", bad, b"\x00" * 4)
@pytest.mark.parametrize("changes", [
{"encoding": "pickle"}, {"encoding": None}, {"encoding": []},
{"dtype": "uint8"}, {"dtype": torch.float32}, {"dtype": []},
{"layout": "BHWC"}, {"layout": None}, {"layout": []},
{"shape": (1, 1, 1)}, {"shape": None}, {"shape": "1,1,1"},
{"shape": [1, 1]}, {"shape": [1, 1, 1, 1]}, {"shape": [1, 0, 1]},
{"shape": [1, -1, 1]}, {"shape": [1, True, 1]}, {"shape": [1, 1.0, 1]},
{"shape": [1, "1", 1]}, {"shape": [1, [], 1]}, {"kind": "MASK"},
])
def test_invalid_tensor_descriptor_fields(changes):
with pytest.raises(ValueError):
codec.decode_payload("MASK", descriptor(**changes), b"\x00" * 4)
@pytest.mark.parametrize("metadata", [{}, {"encoding": "tensor-le-v1"}, {"encoding": "utf8-v1", "kind": "STRING"}, {"encoding": []}])
def test_invalid_string_descriptor(metadata):
with pytest.raises(ValueError):
codec.decode_payload("STRING", metadata, b"")
@pytest.mark.parametrize("payload", [b"", b"\x00" * 3, b"\x00" * 5, b"\x00" * 8])
def test_shape_length_mismatch(payload):
with pytest.raises(ValueError, match="length"):
codec.decode_payload("MASK", descriptor(), payload)
@pytest.mark.parametrize("payload", [None, "text", [0, 0, 0, 0], 4])
def test_nonbyte_payload_rejected(payload):
with pytest.raises(ValueError, match="bytes"):
codec.decode_payload("MASK", descriptor(), payload)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
def test_nonfinite_values_fail_on_encode_and_decode(dtype, value):
tensor = torch.tensor([value], dtype=dtype, device="cpu").reshape(1, 1, 1)
with pytest.raises(ValueError, match="finite"):
codec.encode_payload("MASK", tensor)
fmt = {torch.float16: "e", torch.float32: "f", torch.float64: "d"}
if dtype == torch.bfloat16:
payload = struct.pack("<H", 0x7FC0 if math.isnan(value) else (0x7F80 if value > 0 else 0xFF80))
else:
payload = struct.pack("<" + fmt[dtype], value)
metadata = descriptor(dtype=str(dtype).removeprefix("torch."))
with pytest.raises(ValueError, match="finite"):
codec.decode_payload("MASK", metadata, payload)
def test_invalid_utf8_errors_do_not_echo_content():
with pytest.raises(ValueError) as error:
codec.encode_payload("STRING", "private-value\ud800")
assert "private-value" not in str(error.value)
with pytest.raises(ValueError) as error:
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"private-value\xff")
assert "private-value" not in str(error.value)
@pytest.mark.parametrize("max_bytes", [True, False, -1, 4.0, "4", None])
def test_invalid_byte_budget(max_bytes):
with pytest.raises(ValueError, match="max_bytes"):
codec.encode_payload("STRING", "", max_bytes=max_bytes)
with pytest.raises(ValueError, match="max_bytes"):
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"", max_bytes=max_bytes)
def test_exact_byte_budget_and_zero_budget_for_empty_string():
tensor = torch.ones(1, 1, 1, device="cpu")
metadata, payload = codec.encode_payload("MASK", tensor, max_bytes=4)
assert codec.decode_payload("MASK", metadata, payload, max_bytes=4).item() == 1
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("MASK", tensor, max_bytes=3)
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("MASK", metadata, payload, max_bytes=3)
metadata, payload = codec.encode_payload("STRING", "", max_bytes=0)
assert codec.decode_payload("STRING", metadata, payload, max_bytes=0) == ""
for text in ("a", "é", "🎵"):
byte_count = len(text.encode("utf-8"))
metadata, payload = codec.encode_payload("STRING", text, max_bytes=byte_count)
assert codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count) == text
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("STRING", text, max_bytes=byte_count - 1)
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count - 1)
def test_oversized_shape_rejected_before_tensor_allocation(monkeypatch):
assert codec.MAX_PAYLOAD_BYTES == 256 * 1024 * 1024
def forbidden(*args, **kwargs):
pytest.fail("oversized shape must fail before tensor allocation")
for name in ("frombuffer", "empty", "tensor", "zeros", "isfinite"):
monkeypatch.setattr(torch, name, forbidden)
for shape in ([1, 1, codec.MAX_PAYLOAD_BYTES], [2**100, 2**100, 1]):
with pytest.raises(ValueError, match="limit|budget"):
codec.decode_payload("MASK", descriptor(shape=shape), b"\x00" * 4)
def test_oversized_encode_rejected_before_materialization(monkeypatch):
huge_view = torch.ones(1, device="cpu").expand(1, 1, codec.MAX_PAYLOAD_BYTES)
def forbidden(*args, **kwargs):
pytest.fail("oversized tensor must fail before materialization or scanning")
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
monkeypatch.setattr(torch, "isfinite", forbidden)
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_payload("MASK", huge_view)
@pytest.mark.parametrize("itemsize", [2, 4, 8])
def test_native_little_endian_conversion_handles_big_endian(monkeypatch, itemsize):
raw = bytes(range(itemsize * 3))
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="little"))
assert codec._native_little_endian(raw, itemsize) == raw
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="big"))
expected = b"".join(raw[i:i + itemsize][::-1] for i in range(0, len(raw), itemsize))
assert codec._native_little_endian(raw, itemsize) == expected
assert codec._native_little_endian(expected, itemsize) == raw
+314
View File
@@ -0,0 +1,314 @@
"""CPU-only typed collection and optional native V3 mapping acceptance."""
import asyncio
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys
import types
import pytest
import torch
ROOT = Path(__file__).resolve().parents[1]
SOCKETS = ('images', 'audio', 'masks', 'texts', 'video')
WORKERS = ['serverless:pool:1', 'serverless:pool:2']
class Manager:
def __init__(self, *, pending=False, owned=True):
self.pending, self.owned = pending, owned
self.cancelled = []
def owns(self, job, worker):
return self.owned
def is_pending(self, job, worker):
return self.pending
async def cancel(self, **kwargs):
self.cancelled.append(kwargs)
@pytest.fixture
def module(monkeypatch):
package = '_media_collector_test'
for name, path in ((package, ROOT), (package + '.nodes', ROOT / 'nodes'),
(package + '.utils', ROOT / 'utils'),
(package + '.serverless_bridge', ROOT / 'serverless_bridge')):
stub = types.ModuleType(name)
stub.__path__ = [str(path)]
monkeypatch.setitem(sys.modules, name, stub)
io = types.SimpleNamespace(ComfyNode=object, NodeOutput=lambda *items: types.SimpleNamespace(result=items))
api = types.ModuleType('comfy_api.v0_0_2')
api.io = io
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', api)
management = types.ModuleType('comfy.model_management')
management.throw_exception_if_processing_interrupted = lambda: None
comfy = types.ModuleType('comfy')
comfy.model_management = management
monkeypatch.setitem(sys.modules, 'comfy', comfy)
monkeypatch.setitem(sys.modules, 'comfy.model_management', management)
ps = types.SimpleNamespace(distributed_pending_jobs={}, distributed_jobs_lock=asyncio.Lock())
server = types.ModuleType('server')
server.PromptServer = types.SimpleNamespace(instance=ps)
monkeypatch.setitem(sys.modules, 'server', server)
helper = types.ModuleType(package + '.utils.async_helpers')
helper.run_async_in_server_loop = asyncio.run
monkeypatch.setitem(sys.modules, helper.__name__, helper)
spec = importlib.util.spec_from_file_location(package + '.nodes.media_collector', ROOT / 'nodes/media_collector.py')
loaded = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, loaded)
spec.loader.exec_module(loaded)
monkeypatch.setattr(loaded, 'POLL_SECONDS', 0.005)
monkeypatch.setattr(loaded, 'START_GRACE_SECONDS', 0.025)
loaded.test_server = ps
loaded.test_management = management
return loaded
def native_media(offset=0):
return {
'images': [torch.full((1, 2 + offset, 3, 4), offset, dtype=torch.float64),
torch.zeros((2, 1, 1, 3), dtype=torch.float16)],
'audio': [{'waveform': torch.zeros((1, 1 + offset, 8), dtype=torch.float64),
'sample_rate': 44100 + offset}],
'masks': [torch.zeros((1, 2 + offset, 3), dtype=torch.float64)],
'texts': ['', ' café\n🎥\r\n'],
'video': [object()],
}
def setup(module, messages=(), *, manager=None):
queue = asyncio.Queue()
for message in messages:
queue.put_nowait(message)
ps = module.test_server
ps.distributed_pending_jobs = {'job': queue, 'unrelated': asyncio.Queue()}
ps.distributed_serverless = manager if manager is not None else Manager()
return ps, queue
def message(worker, media):
return {'worker_id': worker, 'media': media, 'is_last': True}
def collect(module, **kwargs):
return module.DistributedMediaCollector.execute(**kwargs).result
def remote(module, **kwargs):
controls = dict(multi_job_id=['job'], enabled_worker_ids=[json.dumps(WORKERS)],
expected_media=[json.dumps(['texts'])], delegate_only=[True])
controls.update(kwargs)
return collect(module, **controls)
@pytest.mark.parametrize('controls', [{}, {'pass_through': [True]}, {'enabled_worker_ids': ['[]']}])
def test_all_five_passthrough_preserves_lists_and_native_values(module, controls):
values = native_media()
result = collect(module, **values, **controls)
assert len(result) == 5
for index, socket in enumerate(SOCKETS):
assert result[index] is values[socket]
def test_reversed_arrivals_master_then_workers_then_items(module):
master, first, second = native_media(), native_media(1), native_media(2)
ps, _ = setup(module, [message(WORKERS[1], second), message(WORKERS[0], first)])
result = remote(module, **master, delegate_only=[False], expected_media=[json.dumps(list(SOCKETS))])
for index, socket in enumerate(SOCKETS):
expected = master[socket] + first[socket] + second[socket]
assert len(result[index]) == len(expected)
assert all(actual is original for actual, original in zip(result[index], expected))
assert list(ps.distributed_pending_jobs) == ['unrelated']
assert ps.distributed_serverless.cancelled == []
def test_delegate_only_and_empty_string(module):
ps, _ = setup(module, [message(WORKERS[1], {'texts': ['b']}), message(WORKERS[0], {'texts': ['']})])
result = remote(module)
assert result == ([], [], [], ['', 'b'], [])
assert 'job' not in ps.distributed_pending_jobs
@pytest.mark.parametrize('kwargs', [{}, {'texts': []}, {'texts': ''}, {'texts': ['x'], 'is_worker': [True]},
{'texts': ['x'], 'enabled_worker_ids': ['["direct-worker"]']},
{'texts': ['x'], 'expected_media': ['["images"]']},
{'texts': ['x'], 'pass_through': [True, False]}, {'delegate_only': True},
{'texts': ['x'], 'enabled_worker_ids': ['["serverless:pool:1","serverless:pool:1"]']},
{'texts': ['x'], 'load_balance': ['false']}])
def test_rejects_missing_media_or_invalid_controls(module, kwargs):
with pytest.raises((ValueError, RuntimeError)):
collect(module, **kwargs)
@pytest.mark.parametrize('bad', [
{'worker_id': 'serverless:other:1', 'media': {'texts': ['x']}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': ['x']}, 'is_last': False},
{'worker_id': WORKERS[0], 'media': {'images': [object()]}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': []}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': 'x'}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': [None]}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': ['x'] * 65}, 'is_last': True},
{'worker_id': WORKERS[0], 'media': {'texts': ['x'], 'unknown': ['x']}, 'is_last': True},
{'worker_id': WORKERS[0], 'serverless_error': 'sanitized failure', 'is_last': True},
])
def test_rejects_bad_delivery_without_partial_output_and_cleans(module, bad):
ps, _ = setup(module, [bad])
with pytest.raises((ValueError, RuntimeError)):
remote(module)
assert list(ps.distributed_pending_jobs) == ['unrelated']
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
@pytest.mark.parametrize('after_complete', [False, True])
def test_duplicate_delivery_rejected_including_final_queue_drain(module, after_complete):
first = message(WORKERS[0], {'texts': ['a']})
second = message(WORKERS[1], {'texts': ['b']})
messages = [first, second, first] if after_complete else [first, first, second]
ps, _ = setup(module, messages)
with pytest.raises(RuntimeError, match='[Dd]uplicate'):
remote(module)
assert 'job' not in ps.distributed_pending_jobs
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
@pytest.mark.parametrize('manager', [None, Manager(owned=False), Manager(pending=False)])
def test_missing_manager_or_finished_participant_fails_bounded(module, manager):
ps, _ = setup(module)
ps.distributed_serverless = manager
with pytest.raises(RuntimeError, match='[Ss]erverless'):
remote(module)
assert 'job' not in ps.distributed_pending_jobs
def test_control_error_also_cleans_exact_job(module):
ps, _ = setup(module)
with pytest.raises(ValueError):
remote(module, expected_media='["unknown"]')
assert list(ps.distributed_pending_jobs) == ['unrelated']
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
def test_launch_race_and_pending_deadline_owned_by_manager(module, monkeypatch):
ps, queue = setup(module)
ps.distributed_serverless = None
async def runner(coro):
async def produce():
await asyncio.sleep(0.01)
ps.distributed_serverless = Manager(pending=True)
await asyncio.sleep(0.035)
queue.put_nowait(message(WORKERS[1], {'texts': ['b']}))
queue.put_nowait(message(WORKERS[0], {'texts': ['a']}))
task = asyncio.create_task(produce())
result = await coro
await task
return result
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
assert remote(module)[3] == ['a', 'b']
assert ps.distributed_serverless.cancelled == []
@pytest.mark.parametrize('cancel_task', [False, True])
def test_cancellation_while_blocked_cleans_and_propagates(module, monkeypatch, cancel_task):
ps, _ = setup(module, manager=Manager(pending=True))
class Interrupted(Exception):
pass
async def runner(coro):
task = asyncio.create_task(coro)
await asyncio.sleep(0.012)
if cancel_task:
task.cancel()
else:
def interrupt():
raise Interrupted('stop')
module.test_management.throw_exception_if_processing_interrupted = interrupt
return await task
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
with pytest.raises(asyncio.CancelledError if cancel_task else Interrupted):
remote(module)
assert list(ps.distributed_pending_jobs) == ['unrelated']
assert ps.distributed_serverless.cancelled == [{'collector_job_id': 'job'}]
def test_replaced_queue_not_removed(module, monkeypatch):
ps, _ = setup(module, manager=Manager(pending=True))
replacement = asyncio.Queue()
async def runner(coro):
task = asyncio.create_task(coro)
await asyncio.sleep(0.012)
ps.distributed_pending_jobs['job'] = replacement
return await task
monkeypatch.setattr(module, 'run_async_in_server_loop', lambda coro: asyncio.run(runner(coro)))
with pytest.raises(RuntimeError, match='queue'):
remote(module)
assert ps.distributed_pending_jobs['job'] is replacement
def test_native_v3_schema_and_list_mapping_cpu():
core = os.environ.get('COMFYUI_SOURCE_ROOT')
if not core:
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only native V3 mapping')
probe = r'''
import asyncio, importlib.util, sys, types
from pathlib import Path
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
import execution
import torch
import server
# No HTTP listener or PromptServer bootstrap is needed for schema/list mapping.
server.PromptServer.instance = types.SimpleNamespace()
root = Path(sys.argv[2])
for name, path in [('typed_collector', root), ('typed_collector.nodes', root/'nodes'),
('typed_collector.utils', root/'utils'),
('typed_collector.serverless_bridge', root/'serverless_bridge')]:
package = types.ModuleType(name)
package.__path__ = [str(path)]
sys.modules[name] = package
spec = importlib.util.spec_from_file_location('typed_collector.nodes.media_collector', root/'nodes/media_collector.py')
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
cls = module.DistributedMediaCollector
info = cls.GET_NODE_INFO_V1()
assert info['output'] == ['IMAGE','AUDIO','MASK','STRING','VIDEO'], info
assert info['output_is_list'] == [True]*5, info
assert cls.define_schema().is_input_list and cls.define_schema().accept_all_inputs
for name, kind in zip(['images','audio','masks','texts','video'], info['output']):
field = info['input']['optional'][name]
assert field[0] == kind and field[1]['forceInput'] is True, field
assert info['input']['required']['load_balance'][1]['default'] is False
async def check():
values = {'load_balance':False, 'pass_through':True}
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe')
assert not missing, missing
image = torch.zeros((1,2,3,4), dtype=torch.float64)
audio = {'waveform':torch.zeros((1,2,7)), 'sample_rate':48000}
mask = torch.ones((1,2,3), dtype=torch.float64)
from comfy_api.input_impl import VideoFromComponents
from comfy_api.util import VideoComponents
from fractions import Fraction
video = VideoFromComponents(VideoComponents(images=image[:,:,:,:3].float(), frame_rate=Fraction(24,1)))
native = {'images':[image], 'audio':[audio], 'masks':[mask], 'texts':[''], 'video':[video]}
prepared.update(native)
results = await execution._async_map_node_over_list('cpu','probe',cls,prepared,cls.FUNCTION,v3_data=hidden)
results = await execution.resolve_map_node_over_list_results(results)
assert len(results) == 1
for index, name in enumerate(native):
assert results[0].result[index] is native[name]
assert results[0].result[4][0] is video
print('MEDIA_COLLECTOR_CPU_OK')
asyncio.run(check())
'''
env = {**os.environ, 'CUDA_VISIBLE_DEVICES': '', 'HIP_VISIBLE_DEVICES': '', 'PYTHONDONTWRITEBYTECODE': '1'}
env.pop('PYTHONPATH', None)
result = subprocess.run([sys.executable, '-c', probe, core, str(ROOT)], env=env,
capture_output=True, text=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
assert 'MEDIA_COLLECTOR_CPU_OK' in result.stdout
+152
View File
@@ -0,0 +1,152 @@
"""Offline contract tests: malformed results never reach transport/delivery."""
import asyncio
import copy
import importlib.util
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location("media_contract_test", ROOT / "serverless_bridge" / "contract.py")
media = importlib.util.module_from_spec(spec)
spec.loader.exec_module(media)
def binding(**changes):
return dict(binding_id="binding-1", collector_id="100", job_id="job-1",
participant_id="serverless:worker:0", node_id="export-1", socket="texts", kind="STRING", **changes)
def response():
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b" \n")
return {"version": 1, "profile": media.PROFILE, "request_id": "request-1",
"bindings": [dict(binding(), count=1)], "artifacts": [artifact],
"items": [{"binding_id": "binding-1", "item_index": 0, "artifact_id": "a"}]}
def validate(value):
return media.validate_response(value, request_id="request-1", bindings=[binding()])
def test_valid_manifest_and_exact_inline_payload():
value = response()
ordered = validate(value)
assert asyncio.run(media.resolve_artifact(ordered["binding-1"][0])) == b" \n"
def test_order_uses_indices_not_arrival():
value = response()
value["bindings"][0]["count"] = 2
value["artifacts"].append(media.artifact_from_payload("b", "STRING", {"encoding": "utf8-v1"}, b"second"))
value["items"] = [{"binding_id": "binding-1", "item_index": 1, "artifact_id": "b"}, value["items"][0]]
assert [item["id"] for item in validate(value)["binding-1"]] == ["a", "b"]
@pytest.mark.parametrize("field,value", [("version", True), ("version", 2), ("profile", "legacy"), ("request_id", "another")])
def test_envelope_identity(field, value):
obj = response()
obj[field] = value
with pytest.raises(ValueError):
validate(obj)
@pytest.mark.parametrize("field", ["collector_id", "participant_id", "job_id", "node_id", "binding_id"])
def test_binding_identity(field):
value = response()
value["bindings"][0][field] = "another"
with pytest.raises(ValueError):
validate(value)
@pytest.mark.parametrize("count", [0, -1, True, 1.0, "1", 65])
def test_invalid_count(count):
value = response()
value["bindings"][0]["count"] = count
with pytest.raises(ValueError):
validate(value)
@pytest.mark.parametrize("mutation", ["missing_binding", "missing_item", "extra_item", "duplicate_binding", "duplicate_artifact", "unused_artifact", "wrong_type", "wrong_index"])
def test_incomplete_or_ambiguous_manifest(mutation):
value = response()
if mutation == "missing_binding":
value["bindings"] = []
elif mutation == "missing_item":
value["items"] = []
elif mutation == "extra_item":
value["items"] *= 2
elif mutation == "duplicate_binding":
value["bindings"] *= 2
elif mutation == "duplicate_artifact":
value["artifacts"] *= 2
elif mutation == "unused_artifact":
artifact = copy.deepcopy(value["artifacts"][0])
artifact["id"] = "unused"
value["artifacts"].append(artifact)
elif mutation == "wrong_type":
value["artifacts"][0]["kind"] = "MASK"
elif mutation == "wrong_index":
value["items"][0]["item_index"] = True
with pytest.raises(ValueError):
validate(value)
def test_duplicate_item_indices():
value = response()
value["bindings"][0]["count"] = 2
value["items"] *= 2
with pytest.raises(ValueError, match="Duplicate"):
validate(value)
def test_total_budget():
with pytest.raises(ValueError, match="total byte"):
media.validate_response(response(), request_id="request-1", bindings=[binding()], max_total_bytes=1)
def test_checksum_and_base64_validation():
artifact = response()["artifacts"][0]
artifact["sha256"] = "0" * 64
with pytest.raises(ValueError, match="checksum"):
asyncio.run(media.resolve_artifact(artifact))
artifact["transport"]["data_base64"] = "!!!!"
with pytest.raises(ValueError, match="base64"):
asyncio.run(media.resolve_artifact(artifact))
def test_empty_text_is_a_value():
artifact = media.artifact_from_payload("empty", "STRING", {"encoding": "utf8-v1"}, b"")
assert asyncio.run(media.resolve_artifact(artifact)) == b""
@pytest.mark.parametrize("url,expires", [("http://example.com/a", "2099-01-01T00:00:00Z"),
("https://user:secret@example.com/a", "2099-01-01T00:00:00Z"),
("https://example.com/a", "2000-01-01T00:00:00Z"),
("https://example.com/a", "2099-01-01"),
("https://example.com:1234/a", "2099-01-01T00:00:00Z")])
def test_invalid_url_envelopes(url, expires):
with pytest.raises(ValueError):
media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"a",
transport={"type": "url", "url": url, "expires_at": expires})
def test_url_resolution_is_bounded_and_hash_verified():
artifact = media.artifact_from_payload("a", "STRING", {"encoding": "utf8-v1"}, b"hello",
transport={"type": "url", "url": "https://example.com/a", "expires_at": "2099-01-01T00:00:00Z"})
calls = []
async def download(url, *, max_bytes):
calls.append((url, max_bytes))
return b"hello"
assert asyncio.run(media.resolve_artifact(artifact, download=download)) == b"hello"
assert calls == [("https://example.com/a", 5)]
def test_descriptor_depth_and_nonfinite_rejected():
value = {}
cursor = value
for _ in range(15):
cursor["next"] = {}
cursor = cursor["next"]
for descriptor in [value, {"x": float("nan")}, {"x": object()}]:
with pytest.raises(ValueError):
media.artifact_from_payload("a", "STRING", descriptor, b"")
+47
View File
@@ -0,0 +1,47 @@
"""Opt-in, isolated CPU acceptance of the complete typed-media data path.
COMFYUI_SOURCE_ROOT must name an existing compatible checkout; run with that
checkout's Python environment. No installation, normal instance, or provider is
used. The child owns its loopback server and every writable Comfy directory.
"""
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import pytest
ROOT = Path(__file__).resolve().parents[1]
def test_typed_media_end_to_end_cpu():
configured = os.environ.get("COMFYUI_SOURCE_ROOT")
if configured is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for real CPU typed-media end-to-end acceptance")
core = Path(configured).expanduser().resolve()
assert configured and all((core / name).is_file() for name in (
"execution.py", "server.py", "comfy/cli_args.py",
"comfy_api/latest/_input_impl/video_types.py",
)), "Configured COMFYUI_SOURCE_ROOT must provide PromptExecutor and native V3 VIDEO"
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env = dict(os.environ)
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
env.pop(name, None)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="", ROCR_VISIBLE_DEVICES="",
PYTHONDONTWRITEBYTECODE="1", PYTHONHASHSEED="0", COMFYUI_IS_WORKER="1",
OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
with tempfile.TemporaryDirectory(prefix="media-e2e-", dir=scratch) as run_root:
env["TMPDIR"] = run_root
result = subprocess.run(
[sys.executable, str(ROOT / "tests/helpers/media_end_to_end_check.py"),
str(core), run_root], cwd=run_root, env=env,
capture_output=True, text=True, timeout=150,
)
assert result.returncode == 0, result.stdout + result.stderr
for marker in ("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK", "REAL_HTTP_BRIDGE_RESULTS_OK",
"MASTER_AND_DELEGATE_ORDER_OK", "ATOMIC_DECODE_REJECTION_OK",
"MEDIA_END_TO_END_CPU_OK"):
assert marker in result.stdout, result.stdout + result.stderr
print(result.stdout)
+877
View File
@@ -0,0 +1,877 @@
"""CPU-only endpoint adapter tests against an actual loopback Comfy HTTP fake."""
import asyncio
import base64
import copy
import hashlib
import importlib
import json
import sys
import types
from contextlib import asynccontextmanager
from pathlib import Path
import aiohttp
from aiohttp import web
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "media_handler_testpkg"
package = types.ModuleType(PACKAGE)
package.__path__ = [str(ROOT / "serverless_bridge")]
sys.modules[PACKAGE] = package
handler = importlib.import_module(PACKAGE + ".handler")
contract = importlib.import_module(PACKAGE + ".contract")
def request(kind="STRING", *, incoming=True):
socket = {value: key for key, value in contract.SOCKET_KINDS.items()}[kind]
binding = dict(binding_id="binding:1", collector_id="collector:1", job_id="job/42",
participant_id="worker:0", node_id="export:1", socket=socket, kind=kind)
artifact = contract.artifact_from_payload("input:1", kind, {"encoding": "test"}, b"input bytes")
return {"workflow": {
"source:1": {"class_type": "Source", "inputs": {"value": "literal"}},
"export:1": {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(binding), socket: ["source:1", 0]}},
}, "media": {"version": 1, "profile": contract.PROFILE, "request_id": "request:1",
"artifacts": [artifact] if incoming else [],
"inputs": [{"node_id": "source:1", "input_name": "value", "artifact_id": "input:1"}] if incoming else [],
"outputs": [binding]}}
class Comfy:
def __init__(self, root, data=None, *, payload=b"result bytes"):
self.root = root
self.data = data or request()
self.payload = payload
self.submitted = []
self.views = []
self.deletes = []
self.interrupts = []
self.polls = 0
self.active = False
self.mode = "success"
self.history_failures = 0
self.record_mutation = None
self.prompt_seen = asyncio.Event()
self.terminal_seen = False
def staged(self):
return list(self.root.glob("distributed-media-v1/*/*.bin"))
def record(self):
binding = self.data["media"]["outputs"][0]
aid = "a" * 32
artifact = contract.artifact_from_payload(aid, binding["kind"], {"encoding": "test"}, self.payload)
artifact["transport"] = {"type": "comfy-file", "filename": aid + ".bin",
"subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(),
"folder_type": "temp"}
value = {"request_id": "request:1", "binding": dict(binding, count=1), "artifacts": [artifact],
"items": [{"binding_id": binding["binding_id"], "item_index": 0, "artifact_id": aid}]}
if self.record_mutation:
self.record_mutation(value)
return value
async def submit(self, req):
data = await req.json()
self.submitted.append(data)
self.active = True
self.prompt_seen.set()
if self.data["media"]["artifacts"]:
assert self.staged()
link = data["prompt"]["source:1"]["inputs"]["value"]
imported = data["prompt"][link[0]]
assert imported["class_type"] == "MediaImportV1"
metadata = json.loads(imported["inputs"]["artifact_json"])
assert "transport" not in metadata
assert (self.root / imported["inputs"]["filename"]).read_bytes() == b"input bytes"
if self.mode == "ambiguous":
return web.Response(status=500)
if self.mode == "partial":
return web.json_response({"prompt_id": "prompt:1", "node_errors": {"export:1": {"errors": ["bad"]}}})
return web.json_response({"prompt_id": "prompt:1", "node_errors": {}})
async def history(self, req):
self.polls += 1
if self.history_failures:
self.history_failures -= 1
return web.Response(status=503)
if self.mode in {"queued", "running", "unrelated", "ambiguous"}:
return web.json_response({})
if self.data["media"]["artifacts"]:
assert self.staged(), "staging must survive until terminal confirmation"
self.active = False
self.terminal_seen = True
if self.mode in {"error", "interrupted"}:
status = {"completed": False, "status_str": "error", "messages": [["execution_interrupted", {}]]}
else:
status = {"completed": True, "status_str": "success", "messages": []}
return web.json_response({"prompt:1": {"status": status, "outputs": {} if self.mode == "missing" else {
"export:1": {"media_exports": [self.record()]}}}})
async def view(self, req):
self.views.append(dict(req.query))
assert self.terminal_seen
payload = self.payload if self.mode != "corrupt" else b"X" * len(self.payload)
if self.mode == "oversized":
payload += b"X"
return web.Response(body=payload)
def entry(self, pid="prompt:1"):
extra = dict(self.submitted[0].get("extra_data", {}), client_id=self.submitted[0]["client_id"])
return [1, pid, self.submitted[0]["prompt"], extra, ["export:1"]]
async def queue(self, req):
if req.method == "POST":
body = await req.json()
assert body == {"delete": ["prompt:1"]}
self.deletes.append(body)
assert self.staged()
self.mode = "deleted"
self.active = False
return web.Response(status=200)
running = [self.entry()] if self.mode == "running" else []
if self.mode == "unrelated":
running = [self.entry("not-our-prompt")]
pending = [self.entry()] if self.mode in {"queued", "ambiguous"} else []
return web.json_response({"queue_running": running, "queue_pending": pending})
async def interrupt(self, req):
assert self.mode == "running"
assert self.staged()
self.interrupts.append(await req.json())
self.mode = "interrupted"
return web.Response(status=200)
@asynccontextmanager
async def serve(comfy):
app = web.Application()
app.add_routes([web.post("/prompt", comfy.submit), web.get("/history/{prompt_id}", comfy.history),
web.get("/view", comfy.view), web.get("/queue", comfy.queue),
web.post("/queue", comfy.queue), web.post("/interrupt", comfy.interrupt)])
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", 0)
await site.start()
port = site._server.sockets[0].getsockname()[1]
try:
yield f"http://127.0.0.1:{port}"
finally:
await runner.cleanup()
def run(coro):
return asyncio.run(coro)
@pytest.mark.parametrize("kind,index", [("IMAGE", 0), ("AUDIO", 1), ("MASK", 2), ("STRING", 3), ("VIDEO", 4)])
def test_loopback_submit_poll_view_cleanup_and_caller_ownership(tmp_path, kind, index):
async def check():
data = request(kind)
original = copy.deepcopy(data)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url, aiohttp.ClientSession() as session:
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, session=session)
assert not session.closed
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
assert await contract.resolve_artifact(result["media"]["artifacts"][0]) == comfy.payload
assert len(comfy.submitted) == 1
assert comfy.submitted[0]["partial_execution_targets"] == ["export:1"]
assert comfy.submitted[0]["prompt"]["source:1"]["inputs"]["value"][1] == index
assert comfy.views == [{"filename": "a" * 32 + ".bin", "subfolder": "distributed-media-v1/" + hashlib.sha256(b"request:1").hexdigest(), "type": "temp"}]
assert data == original
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mode,match", [("corrupt", "checksum"), ("missing", "Missing"), ("error", "execution"), ("oversized", "length")])
def test_terminal_failures_clean_only_after_terminal(tmp_path, mode, match):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises((ValueError, RuntimeError), match=match):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mutation", [
lambda d: d["media"].update(version=True),
lambda d: d["media"].update(profile="legacy"),
lambda d: d["media"].update(request_id="../escape"),
lambda d: d["media"]["outputs"][0].update(node_id="missing"),
lambda d: d["workflow"]["export:1"].update(class_type="SaveImage"),
lambda d: d["workflow"]["export:1"]["inputs"].update(request_id="other"),
lambda d: d["workflow"]["export:1"]["inputs"].update(binding_json="{}"),
lambda d: d["workflow"]["export:1"]["inputs"].update(images=["source:1", 0]),
lambda d: d["workflow"]["export:1"]["inputs"].update(texts="literal"),
lambda d: d["workflow"].update(unrelated={"class_type": "SaveImage", "inputs": {}}),
lambda d: d["workflow"].update(unsafe={"class_type": "MediaImportV1", "inputs": {}}),
lambda d: d["media"]["inputs"].append(copy.deepcopy(d["media"]["inputs"][0])),
lambda d: d["media"]["inputs"][0].update(node_id="missing"),
lambda d: d["media"]["inputs"][0].update(input_name="missing"),
lambda d: d["media"].update(inputs=[]),
lambda d: d["media"]["artifacts"][0].update(sha256="0" * 64),
])
def test_rejection_before_submit(tmp_path, mutation):
async def check():
data = request()
mutation(data)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("mutation", [
lambda r: r.update(request_id="other"),
lambda r: r["binding"].update(count=2),
lambda r: r["binding"].update(collector_id="other"),
lambda r: r["items"][0].update(item_index=1),
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
lambda r: r["artifacts"][0]["transport"].update(folder_type="input"),
lambda r: r["artifacts"][0].update(byte_length=contract.MAX_ARTIFACT_BYTES + 1),
])
def test_bad_result_manifest_rejected_before_any_file_read(tmp_path, mutation):
async def check():
comfy = Comfy(tmp_path)
comfy.record_mutation = mutation
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
run(check())
def test_idempotent_history_retry_without_resubmission(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.history_failures = 2
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert comfy.polls == 3
assert len(comfy.submitted) == 1
run(check())
@pytest.mark.parametrize("mode", ["queued"])
def test_deadline_cancels_only_own_prompt_before_cleanup(tmp_path, monkeypatch, mode):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
async def check():
comfy = Comfy(tmp_path)
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises(TimeoutError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
assert len(comfy.submitted) == 1
assert bool(comfy.deletes) == (mode == "queued")
assert bool(comfy.interrupts) == (mode == "running")
assert not comfy.staged()
run(check())
def test_unconfirmed_terminal_retains_files_and_never_interrupts_unrelated(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.06)
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "unrelated"
async with serve(comfy) as url:
with pytest.raises(RuntimeError, match="retained"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, timeout_seconds=0.04)
assert comfy.staged()
assert not comfy.interrupts
assert not comfy.deletes
assert len(comfy.submitted) == 1
run(check())
def test_ambiguous_post_is_not_retried_and_own_queue_entry_is_reconciled(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "ambiguous"
async with serve(comfy) as url:
with pytest.raises(RuntimeError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert len(comfy.deletes) == 1
assert not comfy.staged()
run(check())
def test_task_cancellation_fails_closed_and_preserves_caller_session(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "running"
async with serve(comfy) as url, aiohttp.ClientSession() as session:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, session=session))
await comfy.prompt_seen.wait()
await asyncio.sleep(0.01)
task.cancel()
with pytest.raises(handler.MediaExecutionUncertain, match="prompt-scoped interruption unavailable"):
await task
assert not session.closed
assert not comfy.interrupts
assert comfy.staged()
run(check())
def test_legacy_images_are_owned_and_rewritten_without_overwriting_user_files(tmp_path):
async def check():
data = request(incoming=False)
data["workflow"]["source:1"] = {"class_type": "LoadImage", "inputs": {"image": "same.png"}}
data["images"] = [{"name": "same.png", "image": base64.b64encode(b"new image").decode()}]
(tmp_path / "same.png").write_bytes(b"user image")
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
filename = comfy.submitted[0]["prompt"]["source:1"]["inputs"]["image"]
assert filename.startswith("distributed-media-v1/")
assert filename.endswith(".bin")
assert (tmp_path / "same.png").read_bytes() == b"user image"
assert not comfy.staged()
run(check())
def test_url_input_uses_bounded_injected_downloader(tmp_path):
async def check():
data = request()
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/a", "expires_at": "2099-01-01T00:00:00Z"}
calls = []
async def download(url, *, max_bytes):
calls.append((url, max_bytes))
return b"input bytes"
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download)
assert calls == [("https://example.org/a", len(b"input bytes"))]
run(check())
def test_large_result_requires_publisher_and_validates_returned_transport(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
comfy = Comfy(tmp_path, payload=b"Z" * 100_000)
async with serve(comfy) as url:
with pytest.raises(ValueError, match="publisher"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
published = []
async def publish(metadata, payload):
assert "transport" not in metadata
assert payload == comfy.payload
published.append(metadata)
return {"type": "url", "url": "https://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
async with serve(comfy) as url:
result = await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=publish)
assert len(published) == 1
assert result["media"]["artifacts"][0]["transport"]["type"] == "url"
assert len(json.dumps(result).encode()) <= 100_000
async def bad_publish(metadata, payload):
return {"type": "url", "url": "http://example.org/result", "expires_at": "2099-01-01T00:00:00Z"}
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path, publish=bad_publish)
run(check())
@pytest.mark.parametrize("name", ["../photo.png", "folder/photo.png", "folder\\photo.png", ".", "..", "/photo.png"])
def test_legacy_basename_rejected_without_submit(tmp_path, name):
async def check():
data = request(incoming=False)
data["images"] = [{"name": name, "image": "YQ=="}]
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
run(check())
def test_staging_rejects_symlink_without_writing_or_submit(tmp_path):
async def check():
outside = tmp_path / "unrelated"
outside.mkdir()
(tmp_path / "distributed-media-v1").symlink_to(outside, target_is_directory=True)
comfy = Comfy(tmp_path)
async with serve(comfy) as url:
with pytest.raises(ValueError, match="symlink"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not list(outside.iterdir())
run(check())
def test_cleanup_preserves_other_runs_in_same_request_namespace(tmp_path):
async def check():
directory = tmp_path / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
directory.mkdir(parents=True)
sibling = directory / ("b" * 32 + ".bin")
sibling.write_bytes(b"other run")
comfy = Comfy(tmp_path)
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert sibling.read_bytes() == b"other run"
assert comfy.staged() == [sibling]
run(check())
def test_input_total_limit_checked_before_downloader_or_submit(tmp_path):
async def check():
data = request()
data["media"]["artifacts"] = []
data["media"]["inputs"] = []
for index in range(3):
artifact = contract.artifact_from_payload(f"input:{index}", "STRING", {}, b"",
transport={"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"})
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
data["media"]["artifacts"].append(artifact)
name = f"value:{index}"
data["workflow"]["source:1"]["inputs"][name] = "literal"
data["media"]["inputs"].append({"node_id": "source:1", "input_name": name, "artifact_id": artifact["id"]})
comfy = Comfy(tmp_path, data)
async def forbidden_download(*args, **kwargs):
pytest.fail("oversized request reached downloader")
async with serve(comfy) as url:
with pytest.raises(ValueError, match="total byte"):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=forbidden_download)
assert not comfy.submitted
run(check())
def test_output_total_limit_checked_before_fetch(tmp_path):
def excessive(record):
first = record["artifacts"][0]
record["artifacts"] = []
record["items"] = []
record["binding"]["count"] = 3
for index, char in enumerate("abc"):
artifact = copy.deepcopy(first)
artifact["id"] = char * 32
artifact["byte_length"] = contract.MAX_ARTIFACT_BYTES
artifact["transport"]["filename"] = artifact["id"] + ".bin"
record["artifacts"].append(artifact)
record["items"].append({"binding_id": "binding:1", "item_index": index, "artifact_id": artifact["id"]})
async def check():
comfy = Comfy(tmp_path)
comfy.record_mutation = excessive
async with serve(comfy) as url:
with pytest.raises(ValueError, match="total byte"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert not comfy.views
assert not comfy.staged()
run(check())
def test_partial_target_acceptance_is_not_reported_as_success(tmp_path):
async def check():
comfy = Comfy(tmp_path)
comfy.mode = "partial"
async with serve(comfy) as url:
with pytest.raises(RuntimeError, match="incomplete"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=tmp_path)
assert len(comfy.submitted) == 1
assert not comfy.views
assert not comfy.staged()
run(check())
# Native routes are compiled from the pinned source, without importing ComfyUI
# (which would initialize the GPU). All dependencies and HTTP are CPU-only.
@pytest.mark.parametrize("current,requested,race,expected", [
("own", "own", False, ["own"]),
("other", "own", False, []),
("own", "own", True, ["successor"]),
])
def test_pinned_native_interrupt_route_has_snapshot_race(current, requested, race, expected):
import ast
import logging
source = Path("/home/robert/ComfyUI_dev/server.py")
if not source.is_file():
pytest.skip("Pinned ComfyUI source unavailable")
tree = ast.parse(source.read_text())
routes = [node for node in ast.walk(tree) if isinstance(node, ast.AsyncFunctionDef)
and any("/interrupt" in ast.unparse(d) for d in node.decorator_list)]
assert len(routes) == 1
route = routes[0]
route.decorator_list = []
interrupts = []
class Queue:
running = current
def get_current_queue(self):
snapshot = [[1, self.running, {}, {}, []]]
if race:
self.running = "successor"
return snapshot, []
queue = Queue()
scope = {"json": json, "logging": logging, "web": web,
"self": types.SimpleNamespace(prompt_queue=queue),
"nodes": types.SimpleNamespace(interrupt_processing=lambda: interrupts.append(queue.running))}
exec(compile(ast.Module(body=[route], type_ignores=[]), str(source), "exec"), scope)
async def check():
async def body():
return {"prompt_id": requested}
response = await scope[route.name](types.SimpleNamespace(json=body))
assert response.status == 200
run(check())
assert interrupts == expected
def exporter_file(root, *, name="a" * 32 + ".bin", payload=b"result bytes"):
directory = root / "distributed-media-v1" / hashlib.sha256(b"request:1").hexdigest()
directory.mkdir(parents=True, exist_ok=True)
path = directory / name
path.write_bytes(payload)
return path
@pytest.mark.parametrize("mode", ["success", "corrupt", "oversized", "error", "interrupted", "partial"])
def test_owned_export_cleanup_after_terminal_success_or_failure(tmp_path, mode):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
sibling = exporter_file(temp, name="b" * 32 + ".bin", payload=b"unrelated")
ordinary = exporter_file(temp, name="keep.txt", payload=b"ordinary")
comfy = Comfy(tmp_path / "input")
comfy.mode = mode
view = comfy.view
async def reading(req):
assert owned.exists(), "cleanup raced artifact collection"
return await view(req)
comfy.view = reading
async with serve(comfy) as url:
if mode == "success":
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
else:
with pytest.raises((ValueError, RuntimeError)):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert comfy.terminal_seen
assert not owned.exists()
assert sibling.read_bytes() == b"unrelated"
assert ordinary.read_bytes() == b"ordinary"
assert not comfy.staged()
run(check())
def test_omitted_temp_root_leaves_exporter_lifecycle_ownership(tmp_path):
async def check():
owned = exporter_file(tmp_path / "temp")
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root)
assert owned.read_bytes() == comfy.payload
run(check())
@pytest.mark.parametrize("mode", ["running", "unrelated"])
def test_uncertain_exports_and_inputs_retained_without_interrupt(tmp_path, monkeypatch, mode):
monkeypatch.setattr(handler, "POLL_INTERVAL_SECONDS", 0.005)
monkeypatch.setattr(handler, "CANCEL_GRACE_SECONDS", 0.04)
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.mode = mode
async with serve(comfy) as url:
with pytest.raises(handler.MediaExecutionUncertain):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
temp_dir=temp, timeout_seconds=0.03)
assert owned.exists()
assert comfy.staged()
assert not comfy.interrupts
assert not comfy.deletes
run(check())
@pytest.mark.parametrize("component", ["root", "namespace", "request", "file"])
def test_export_cleanup_never_traverses_symlinks(tmp_path, component):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
selected = {"root": temp, "namespace": owned.parent.parent,
"request": owned.parent, "file": owned}[component]
outside = tmp_path / "outside"
selected.rename(outside)
selected.symlink_to(outside, target_is_directory=component != "file")
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert selected.is_symlink()
assert owned.read_bytes() == comfy.payload
run(check())
@pytest.mark.parametrize("mutation", [
lambda r: r.update(request_id="other"),
lambda r: r["binding"].update(collector_id="other"),
lambda r: r["artifacts"][0]["transport"].update(filename="../escape"),
lambda r: r["artifacts"][0]["transport"].update(subfolder="distributed-media-v1/other"),
])
def test_unidentified_exports_are_not_deleted(tmp_path, mutation):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.record_mutation = mutation
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert not comfy.views
run(check())
@pytest.mark.parametrize("failure", ["required", "publish"])
def test_publication_failures_cleanup_identified_exports(tmp_path, monkeypatch, failure):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
async def publish(metadata, payload):
assert owned.exists()
raise ValueError("publisher failed")
async with serve(comfy) as url:
with pytest.raises(ValueError, match="publisher"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root,
temp_dir=temp, publish=publish if failure == "publish" else None)
assert not owned.exists()
run(check())
def test_repeated_cancel_drains_collection_before_cleanup(tmp_path, monkeypatch):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def publish(metadata, payload):
started.set()
try:
await asyncio.Future()
finally:
draining.set()
await release.wait()
assert owned.exists(), "cleanup while publisher still reading"
async with serve(comfy) as url:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
input_dir=comfy.root, temp_dir=temp, publish=publish))
await asyncio.wait_for(started.wait(), 1)
task.cancel()
await asyncio.wait_for(draining.wait(), 1)
task.cancel()
await asyncio.sleep(0)
assert owned.exists()
release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert not owned.exists()
run(check())
@pytest.mark.parametrize("cancel_caller", [False, True])
def test_stubborn_publisher_teardown_is_bounded_and_retains_files(tmp_path, monkeypatch, cancel_caller):
monkeypatch.setattr(handler, "MAX_RESPONSE_JSON_BYTES", 100_000)
monkeypatch.setattr(handler, "DRAIN_GRACE_SECONDS", 0.03)
async def check():
temp = tmp_path / "temp"
comfy = Comfy(tmp_path / "input", payload=b"Z" * 100_000)
owned = exporter_file(temp, payload=comfy.payload)
started, draining, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def publish(metadata, payload):
started.set()
try:
await asyncio.Future()
finally:
draining.set()
await release.wait()
assert owned.exists()
async with serve(comfy) as url:
task = asyncio.create_task(handler.execute_media_request(comfy.data, comfy_url=url,
input_dir=comfy.root, temp_dir=temp, publish=publish,
timeout_seconds=5 if cancel_caller else 0.15))
try:
await asyncio.wait_for(started.wait(), 1)
if cancel_caller:
task.cancel()
await asyncio.wait_for(draining.wait(), 1)
if cancel_caller:
task.cancel() # must not reset the grace or recancel the reader
done, _ = await asyncio.wait({task}, timeout=0.5)
assert task in done, "callback suppressed cancellation indefinitely"
with pytest.raises(handler.MediaExecutionUncertain, match="teardown"):
task.result()
assert comfy.terminal_seen and owned.exists() and comfy.staged()
assert handler._UNCERTAIN_WORK
pending = list(handler._UNCERTAIN_WORK)
release.set()
await asyncio.gather(*pending, return_exceptions=True)
await asyncio.sleep(0)
assert not handler._UNCERTAIN_WORK
assert owned.exists() and comfy.staged(), "late callback must not retire files"
finally:
release.set()
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
run(check())
def test_export_replaced_during_collection_is_preserved(tmp_path):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
view = comfy.view
async def replace(req):
replacement = owned.with_suffix(".replacement")
replacement.write_bytes(b"unrelated replacement")
replacement.replace(owned)
return await view(req)
comfy.view = replace
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.read_bytes() == b"unrelated replacement"
run(check())
@pytest.mark.parametrize("status", ["success", "error"])
def test_partial_history_cleans_only_valid_identified_exporters(tmp_path, status):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
unidentified = exporter_file(temp, name="c" * 32 + ".bin")
data = request()
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
data["media"]["outputs"].append(second)
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(second), "texts": ["source:1", 0]}}
comfy = Comfy(tmp_path / "input", data)
comfy.mode = status
async with serve(comfy) as url:
with pytest.raises((ValueError, RuntimeError)):
await handler.execute_media_request(data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert not owned.exists()
assert unidentified.exists()
assert not comfy.views
run(check())
def test_missing_history_exports_remain_lifecycle_owned(tmp_path):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
comfy = Comfy(tmp_path / "input")
comfy.mode = "missing"
async with serve(comfy) as url:
with pytest.raises(ValueError, match="Missing"):
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("replacement", ["symlink", "directory", "hardlink"])
def test_nonregular_or_linked_export_files_are_preserved(tmp_path, replacement):
async def check():
temp = tmp_path / "temp"
owned = exporter_file(temp)
outside = tmp_path / "outside"
outside.write_bytes(b"unrelated")
if replacement == "hardlink":
alias = tmp_path / "alias"
alias.hardlink_to(owned)
else:
owned.unlink()
if replacement == "symlink":
owned.symlink_to(outside)
else:
owned.mkdir()
comfy = Comfy(tmp_path / "input")
async with serve(comfy) as url:
await handler.execute_media_request(comfy.data, comfy_url=url, input_dir=comfy.root, temp_dir=temp)
assert owned.exists()
assert outside.read_bytes() == b"unrelated"
run(check())
@pytest.mark.parametrize("shared", [False, True])
def test_direct_exporter_input_prunes_only_obsolete_loader_ancestors(tmp_path, shared):
async def check():
data = request("VIDEO")
data["workflow"]["source:1"] = {"class_type": "LoadVideo", "inputs": {"file": "unavailable.mp4"}}
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": "video", "artifact_id": "input:1"}]
if shared:
second = dict(data["media"]["outputs"][0], binding_id="binding:2", node_id="export:2", collector_id="collector:2")
data["media"]["outputs"].append(second)
data["workflow"]["export:2"] = {"class_type": "MediaExportV1", "inputs": {
"request_id": "request:1", "binding_json": json.dumps(second), "video": ["source:1", 0]}}
fake_data = copy.deepcopy(data)
fake_data["media"].update(artifacts=[], inputs=[])
comfy = Comfy(tmp_path, fake_data)
history = comfy.history
async def full_history(req):
response = await history(req)
body = json.loads(response.text)
if shared:
record = comfy.record()
record["binding"] = dict(second, count=1)
record["artifacts"][0]["id"] = "b" * 32
record["artifacts"][0]["transport"]["filename"] = "b" * 32 + ".bin"
record["items"] = [{"binding_id": second["binding_id"], "item_index": 0, "artifact_id": "b" * 32}]
body["prompt:1"]["outputs"]["export:2"] = {"media_exports": [record]}
return web.json_response(body)
comfy.history = full_history
async with serve(comfy) as url:
result = await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
contract.validate_response(result["media"], request_id="request:1", bindings=data["media"]["outputs"])
submitted = comfy.submitted[0]["prompt"]
imported, slot = submitted["export:1"]["inputs"]["video"]
assert slot == 4
assert submitted[imported]["class_type"] == "MediaImportV1"
assert json.loads(submitted[imported]["inputs"]["artifact_json"])["kind"] == "VIDEO"
assert ("source:1" in submitted) is shared
assert set(submitted) == ({"export:1", imported, "source:1", "export:2"} if shared else {"export:1", imported})
assert not comfy.staged()
run(check())
@pytest.mark.parametrize("target,kind", [("request_id", "VIDEO"), ("binding_json", "VIDEO"), ("video", "STRING")])
def test_direct_exporter_input_rejects_controls_and_wrong_kind(tmp_path, target, kind):
async def check():
data = request("VIDEO")
data["media"]["inputs"] = [{"node_id": "export:1", "input_name": target, "artifact_id": "input:1"}]
data["media"]["artifacts"][0]["kind"] = kind
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(ValueError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path)
assert not comfy.submitted
assert not comfy.staged()
run(check())
def test_slow_download_times_out_before_any_prompt(tmp_path):
async def check():
data = request()
data["media"]["artifacts"][0]["transport"] = {"type": "url", "url": "https://example.org/input", "expires_at": "2099-01-01T00:00:00Z"}
async def download(*args, **kwargs):
await asyncio.sleep(60)
comfy = Comfy(tmp_path, data)
async with serve(comfy) as url:
with pytest.raises(TimeoutError):
await handler.execute_media_request(data, comfy_url=url, input_dir=tmp_path, download=download, timeout_seconds=0.02)
assert not comfy.submitted
assert not comfy.staged()
run(check())
+330
View File
@@ -0,0 +1,330 @@
"""Real CPU-only IMAGE transfers; no ComfyUI/bootstrap/provider imports."""
import importlib
import sys
import types
from pathlib import Path
import numpy as np
import pytest
import torch
from PIL import Image
# Load a private package, not Distributed's root __init__.
ROOT = Path(__file__).resolve().parents[1] / "serverless_bridge"
PACKAGE = "image_codec_tests_bridge"
pkg = types.ModuleType(PACKAGE)
pkg.__path__ = [str(ROOT)]
sys.modules[PACKAGE] = pkg
codec = importlib.import_module(PACKAGE + ".codec")
images = importlib.import_module(PACKAGE + ".image_codec")
FORMATS = ("legacy_png", "png", "jpeg", "webp", "ffv1", "h264", "h265", "av1")
LOSSLESS = ("legacy_png", "png", "ffv1")
@pytest.fixture(autouse=True)
def no_gpu(monkeypatch):
def forbidden(*a, **kw):
pytest.fail("GPU query or initialization is forbidden")
for name in ("is_available", "init", "_lazy_init", "device_count"):
monkeypatch.setattr(torch.cuda, name, forbidden)
def available(fmt):
if fmt in images._VIDEO_CODECS:
av = pytest.importorskip("av")
name = images._VIDEO_CODECS[fmt][0]
if name not in av.codecs_available:
pytest.skip("Software encoder unavailable: " + name)
def sample(channels=3, dtype=torch.float32, shape=(3, 65, 67)):
g = torch.Generator(device="cpu").manual_seed(42)
return torch.rand((*shape, channels), generator=g, device="cpu", dtype=dtype)
def quantized(value):
return torch.from_numpy(np.clip(value.float().numpy() * 255, 0, 255).astype(np.uint8)).float() / 255
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_exact_unchanged(dtype):
value = sample(4, dtype, (2, 3, 5))
value.flatten()[0] = -0.0
expected = codec.encode_payload("IMAGE", value)
assert images.encode_image(value) == expected
result = images.decode_image(*expected)
assert torch.equal(result.view(torch.uint8), value.view(torch.uint8))
@pytest.mark.parametrize("fmt", FORMATS)
@pytest.mark.parametrize("channels", [1, 3, 4])
def test_real_roundtrips_shape_order_quantization_alpha(fmt, channels):
available(fmt)
value = sample(channels)
# Large inter-frame difference also verifies frame ordering.
value[0] *= 0.2
value[2] = value[2] * 0.2 + 0.8
if channels == 4:
value[..., 3] = torch.linspace(-0.1, 1.1, value[..., 3].numel()).reshape(value[..., 3].shape)
descriptor, payload = images.encode_image(value, transfer_format=fmt)
assert descriptor["encoding"] == "image-transfer-v1"
result = codec.decode_payload("IMAGE", descriptor, payload)
assert result.shape == value.shape
assert result.device.type == "cpu" and result.is_contiguous()
color = 3 if channels == 4 else channels
expected = quantized(value[..., :color])
if fmt in LOSSLESS:
assert torch.equal(result[..., :color], expected)
else:
# Random RGB noise loses chroma detail in 4:2:0 even at high quality.
assert (result[..., :color] - expected).abs().mean() < 0.12
assert result[0, ..., :color].mean() < result[2, ..., :color].mean()
if channels == 4:
assert torch.equal(result[..., 3], value[..., 3])
@pytest.mark.parametrize("fmt", ["png", "ffv1", "jpeg", "webp", "h264", "h265", "av1"])
def test_float64_alpha_is_exact(fmt):
available(fmt)
value = sample(4, torch.float64, (1, 65, 67))
descriptor, payload = images.encode_image(value, transfer_format=fmt)
out = images.decode_image(descriptor, payload)
assert out.dtype == torch.float64
assert torch.equal(out[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
def test_lossy_quality_changes_final_payload_size(fmt):
available(fmt)
value = sample(shape=(2, 96, 96))
low = images.encode_image(value, transfer_format=fmt, transfer_quality=15)[1]
high = images.encode_image(value, transfer_format=fmt, transfer_quality=95)[1]
assert len(high) > len(low)
def test_truncation_clipping_grayscale_noncontiguous():
value = torch.tensor([-0.1, 0.5, 1.1, 0.12345], device="cpu").reshape(1, 2, 2, 1).transpose(1, 2)
out = images.decode_image(*images.encode_image(value, transfer_format="png"))
assert torch.equal(out, quantized(value))
assert out[0, 0, 1, 0] == 1
@pytest.mark.parametrize("value", [None, [], torch.ones(2, 3, 3), torch.ones(1, 2, 3, 2), torch.ones(1, 2, 3, 0), torch.ones(1, 2, 3, 3, dtype=torch.uint8), torch.ones(1, 2, 3, 3, device="meta"), torch.ones(1, 2, 3, 3).to_sparse()])
def test_invalid_values(value):
with pytest.raises(ValueError):
images.encode_image(value, transfer_format="png")
@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")])
def test_nonfinite_rejected(bad):
with pytest.raises(ValueError, match="finite"):
images.encode_image(torch.full((1, 1, 1, 3), bad), transfer_format="png")
@pytest.mark.parametrize("options", [{"transfer_format": "PNG"}, {"transfer_format": []}, {"transfer_quality": True}, {"transfer_quality": 0}, {"transfer_quality": 101}, {"transfer_quality": 1.5}, {"max_bytes": True}, {"max_bytes": -1}])
def test_strict_options(options):
with pytest.raises(ValueError):
images.encode_image(sample(shape=(1, 1, 1)), **options)
@pytest.mark.parametrize("change", [{"extra": 1}, {"format": "exact"}, {"quality": True}, {"layout": "BCHW"}, {"shape": [1, 2**60, 1, 3]}, {"shape": [True, 1, 1, 3]}, {"parts": [0]}, {"parts": [True]}, {"parts": [2**60]}, {"alpha": {"extra": 0}}])
def test_bad_descriptors(change):
descriptor, payload = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
descriptor.update(change)
with pytest.raises(ValueError):
images.decode_image(descriptor, payload)
def test_truncated_appended_and_mismatched_parts():
d, p = images.encode_image(sample(shape=(2, 3, 5)), transfer_format="png")
for payload in (p[:-1], p + b"x"):
with pytest.raises(ValueError):
images.decode_image(d, payload)
for change in ({"shape": [1, 3, 5, 3]}, {"parts": [len(p)]}, {"shape": [2, 2, 5, 3]}):
bad = dict(d, **change)
with pytest.raises(ValueError):
images.decode_image(bad, p)
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp", "ffv1", "h264", "h265", "av1"])
def test_payload_and_decoded_budget(fmt):
available(fmt)
value = sample(shape=(1, 65, 67))
d, p = images.encode_image(value, transfer_format=fmt)
with pytest.raises(ValueError, match="limit|budget"):
images.decode_image(d, p, max_bytes=100)
with pytest.raises(ValueError, match="limit|budget"):
images.encode_image(value, transfer_format=fmt, max_bytes=100)
def test_oversized_shape_before_materialization(monkeypatch):
value = torch.ones(1).expand(1, 8192, 8192, 4)
def forbidden(*a, **kw):
pytest.fail("Materialized an oversized tensor")
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
monkeypatch.setattr(torch, "isfinite", forbidden)
with pytest.raises(ValueError):
images.encode_image(value, transfer_format="png")
def test_header_dimensions_checked_before_pixels(monkeypatch):
d, p = images.encode_image(sample(shape=(1, 9, 11)), transfer_format="png")
d["shape"] = [1, 1, 1, 3]
def forbidden(*a, **kw):
pytest.fail("Pixels decoded before header dimensions were checked")
monkeypatch.setattr(Image.Image, "load", forbidden)
with pytest.raises(ValueError):
images.decode_image(d, p)
def test_other_kinds_do_not_dispatch():
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
for kind in ("MASK", "AUDIO", "STRING"):
with pytest.raises(ValueError):
codec.decode_payload(kind, d, p)
def test_missing_video_encoder_is_explicit(monkeypatch):
monkeypatch.setitem(images._VIDEO_CODECS, "h264", ("missing_software_encoder", 51, "yuv420p"))
with pytest.raises(ValueError, match="encoder|unavailable"):
images.encode_image(sample(shape=(1, 64, 64)), transfer_format="h264")
def test_bounded_sink_catches_final_flush_bytes():
sink = images._BoundedIO(4)
sink.write(b"1234")
with pytest.raises(ValueError, match="limit"):
sink.write(b"5")
assert sink.getvalue() == b"1234"
def test_crf_mapping():
assert images._quality_to_crf(90, 51) == 13
assert images._quality_to_crf(90, 63) == 16
assert images._quality_to_crf(100, 63) == 0
@pytest.mark.parametrize("fmt", FORMATS)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_compressed_preserves_source_dtype_and_exact_alpha(fmt, dtype):
available(fmt)
value = sample(4, dtype, (1, 17, 19))
value[..., 3].flatten()[0] = -0.0
d, p = images.encode_image(value, transfer_format=fmt)
with torch.device("meta"):
result = images.decode_image(d, p)
assert result.device.type == "cpu" and result.dtype == dtype
assert d["dtype"] == str(dtype).removeprefix("torch.")
assert torch.equal(result[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
if fmt in LOSSLESS:
assert torch.equal(result[..., :3], quantized(value[..., :3]).to(dtype))
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_count_and_header_fail_before_native_decoder(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
def forbidden(*a, **kw):
pytest.fail("Native decoder reached before header validation")
monkeypatch.setattr(images, "_decode_video", forbidden)
for shape in ([1, 65, 67, 3], [3, 65, 67, 3], [2, 32, 32, 3]):
bad = dict(d, shape=shape)
with pytest.raises(ValueError):
images.decode_image(bad, p)
for payload in (p[:-1], p[:40], p + b"extra"):
bad = dict(d, parts=[len(payload)])
with pytest.raises(ValueError):
images.decode_image(bad, payload)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_does_not_use_native_container_probe(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
av = images._import_av()
def forbidden(*a, **kw):
pytest.fail("Unbounded native stream probing is forbidden")
monkeypatch.setattr(av, "open", forbidden)
assert images.decode_image(d, p).shape == (2, 65, 67, 3)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_real_video_flush_bound(fmt):
available(fmt)
frames = [np.zeros((64, 64, 3), dtype=np.uint8)]
good = images._encode_video(iter(frames), fmt, 90, 64, 64, 100000)
with pytest.raises(ValueError, match="limit"):
images._encode_video(iter(frames), fmt, 90, 64, 64, len(good) - 1)
def test_alpha_invalid_dtype_nonfinite_and_extra_fields():
import copy
d, p = images.encode_image(sample(4, shape=(1, 3, 5)), transfer_format="png")
for field, value in (("dtype", "float64"), ("shape", [1, 3, 6, 1]), ("extra", 1)):
bad = copy.deepcopy(d)
bad["alpha"]["descriptor"][field] = value
with pytest.raises(ValueError):
images.decode_image(bad, p)
import struct
bad = p[:-4] + struct.pack("<f", float("nan"))
with pytest.raises(ValueError, match="finite"):
images.decode_image(d, bad)
def test_descriptor_structure_bound_and_frame_limit():
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
for bad in (dict(d, quality="x" * 70000), dict(d, parts=[1] * 241), dict(d, shape=[241, 1, 1, 3]), dict(d, dtype="x")):
with pytest.raises(ValueError):
images.decode_image(bad, p)
with pytest.raises(ValueError):
images.encode_image(torch.zeros(241, 1, 1, 3), transfer_format="png")
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp"])
def test_all_still_headers_checked_before_pillow_open(fmt, monkeypatch):
d, p = images.encode_image(sample(shape=(1, 31, 33)), transfer_format=fmt)
d["shape"] = [1, 2, 2, 3]
def forbidden(*a, **kw):
pytest.fail("Pillow opened a mismatched header")
monkeypatch.setattr(Image, "open", forbidden)
with pytest.raises(ValueError, match="dimensions"):
images.decode_image(d, p)
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
def test_video_longer_batch_order_and_coded_allocation_bound(fmt):
available(fmt)
value = torch.linspace(0, 1, 16).reshape(16, 1, 1, 1).expand(16, 64, 64, 3)
d, p = images.encode_image(value, transfer_format=fmt)
header = images._video_header(p, fmt, 64, 64, 16)
coded = images._coded_frame_bound(p, fmt, header)
assert 16 <= coded <= images.MAX_FRAMES * 4
result = images.decode_image(d, p)
assert torch.all((result.mean(dim=(1, 2, 3)) - value.mean(dim=(1, 2, 3))).abs() < 0.03)
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
def test_packet_unit_bounds_before_native_decode(fmt):
available(fmt)
d, p = images.encode_image(sample(shape=(1, 64, 64)), transfer_format=fmt)
config, bounds = images._video_header(p, fmt, 64, 64, 1)
start, end = bounds[0]
packet = p[start:end]
with pytest.raises(ValueError):
images._coded_frame_bound(packet[:-1], fmt, (config, [(0, len(packet) - 1)]))
# An adversarial single block can contain many coded frames: bound it too.
enormous = packet * (images.MAX_FRAMES * 4 + 1)
with pytest.raises(ValueError, match="limit"):
images._coded_frame_bound(enormous, fmt, (config, [(0, len(enormous))]))
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
def test_bitstream_resolution_cannot_bypass_native_pixel_limit(fmt, monkeypatch):
available(fmt)
d, p = images.encode_image(sample(shape=(1, 256, 256)), transfer_format=fmt)
header = images._video_header(p, fmt, 256, 256, 1)
av = images._import_av()
# A lying container is not enough to disable max_pixels in the decoder.
# The bitstream still contains 256x256 pixels but the budget permits 64x64.
with pytest.raises((ValueError, av.error.FFmpegError)):
list(images._decode_video(p, fmt, 64, 64, 1, header))
+89
View File
@@ -0,0 +1,89 @@
"""Pure integration of participant identity, typed compilation and delegate pruning."""
import copy
import importlib
import json
import sys
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "media_transform_testpkg"
for suffix in ("", ".api", ".api.orchestration", ".utils", ".serverless_bridge"):
module = types.ModuleType(PACKAGE + suffix)
module.__path__ = [str(ROOT.joinpath(*suffix.strip(".").split(".")))] if suffix else [str(ROOT)]
sys.modules[module.__name__] = module
pt = importlib.import_module(PACKAGE + ".api.orchestration.prompt_transform")
compiler = importlib.import_module(PACKAGE + ".api.orchestration.serverless_media_prompt")
contract = importlib.import_module(PACKAGE + ".serverless_bridge.contract")
def prompt():
nodes = {
"source": {"class_type": "MediaSource", "inputs": {}},
"collector": {"class_type": "DistributedMediaCollector", "inputs": {
socket: ["source", index] for index, socket in enumerate(contract.SOCKET_KINDS)}},
}
for index, socket in enumerate(contract.SOCKET_KINDS):
nodes["sink:" + socket] = {"class_type": "Capture", "inputs": {"value": ["collector", index]}}
return nodes
def prepare(graph, participant, delegate):
index = pt.PromptIndex(graph)
jobs = pt.generate_job_id_map(index, "run")
assert jobs == {"collector": "run_collector"}
return pt.apply_participant_overrides(index.copy_prompt(), participant,
["serverless:pool:1", "serverless:pool:2"], jobs, "http://master", delegate, index)
def test_typed_worker_preparation_compiles_all_bindings_without_source_mutation():
graph = prompt()
before = copy.deepcopy(graph)
worker = prepare(graph, "serverless:pool:2", True)
compiled, bindings = compiler.compile_media_prompt(worker, "request:2")
assert len(bindings) == 5
assert {row["participant_id"] for row in bindings} == {"serverless:pool:2"}
assert {row["job_id"] for row in bindings} == {"run_collector"}
assert {row["socket"] for row in bindings} == set(contract.SOCKET_KINDS)
assert set(compiled) == {"source"} | {row["node_id"] for row in bindings}
assert graph == before
def test_typed_delegate_master_keeps_outputs_without_sources_or_placeholders():
graph = prompt()
prepared = prepare(graph, "master", True)
before = copy.deepcopy(prepared)
delegated = pt.prepare_delegate_master_prompt(prepared, ["collector"])
assert set(delegated) == {"collector"} | {"sink:" + s for s in contract.SOCKET_KINDS}
inputs = delegated["collector"]["inputs"]
assert not set(contract.SOCKET_KINDS) & set(inputs)
assert json.loads(inputs["expected_media"]) == list(contract.SOCKET_KINDS)
assert inputs["delegate_only"] is True
assert prepared == before
def test_text_source_is_not_retained_as_lightweight_delegate_scalar(monkeypatch):
graph = {
"text": {"class_type": "PrimitiveString", "inputs": {"value": "local text"}},
"collector": {"class_type": "DistributedMediaCollector", "inputs": {"texts": ["text", 0]}},
"sink": {"class_type": "Capture", "inputs": {"value": ["collector", 3]}},
}
monkeypatch.setattr(pt, "_delegate_master_input_is_safe_scalar", lambda *args: True)
monkeypatch.setattr(pt, "_is_delegate_master_retained_upstream_node", lambda *args: True)
delegated = pt.prepare_delegate_master_prompt(prepare(graph, "master", True), ["collector"])
assert "text" not in delegated
assert "texts" not in delegated["collector"]["inputs"]
assert json.loads(delegated["collector"]["inputs"]["expected_media"]) == ["texts"]
@pytest.mark.parametrize("socket", contract.SOCKET_KINDS)
def test_participating_master_retains_exact_media_links(socket):
graph = prompt()
original = copy.deepcopy(graph)
master = prepare(graph, "master", False)
assert master["collector"]["inputs"][socket] == graph["collector"]["inputs"][socket]
assert master["collector"]["inputs"]["delegate_only"] is False
assert "worker_id" not in master["collector"]["inputs"]
assert graph == original
+531
View File
@@ -0,0 +1,531 @@
"""Offline queue integration: real compilers, packers and wire validators.
Only host dispatch/probing/queueing and the manager lifecycle are mocked. No
provider requests, services or GPU work are permitted by these tests.
"""
import asyncio
import base64
import copy
import importlib
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import types
import uuid
import wave
from unittest.mock import AsyncMock, Mock, patch
import aiohttp
import pytest
# Preload before the shared integration loader snapshots sys.modules: extension
# modules cannot safely be unloaded and imported again between fixture contexts.
import torch
from PIL import Image
from test_serverless_integration import load_modules, module, PKG
PROFILE = "comfy-media-v1"
def node(class_type, **inputs):
return {"class_type": class_type, "inputs": inputs}
def text_prompt():
return {
"text": node("PrimitiveString", value=" café雪\n"),
"collector": node("DistributedMediaCollector", texts=["text", 0], load_balance=False),
"save": node("TextSink", texts=["collector", 3]),
}
def legacy_prompt():
return {
"source": node("EmptyImage", width=32, height=32, batch_size=1, color=0),
"collector": node("DistributedCollector", images=["source", 0]),
"save": node("SaveImage", images=["collector", 0], filename_prefix="local"),
}
@pytest.fixture
def harness(monkeypatch):
q, modules, ps = load_modules()
with patch.dict(sys.modules, modules):
pool = {"id": "pool", "name": "Media", "type": "serverless", "provider": "runpod",
"endpoint_id": "offline-endpoint", "enabled": True, "jobs_per_run": 2,
"timeout_seconds": 30, "media_profile": PROFILE}
config = {"workers": [pool], "settings": {"master_delegate_only": True}}
q.load_config = lambda: config
q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 7, "node_errors": {}})
manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=Mock())
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: manager)
# These are the actual modules, not permissive serializer/packing stubs.
client = importlib.import_module(PKG + ".api.orchestration.runpod_client")
compiler = importlib.import_module(PKG + ".api.orchestration.serverless_media_prompt")
packer = importlib.import_module(PKG + ".api.orchestration.serverless_typed_inputs")
validator = importlib.import_module(PKG + ".api.orchestration.serverless_typed_results")
handler = importlib.import_module(PKG + ".serverless_bridge.handler")
def forbidden(*args, **kwargs):
raise AssertionError("Offline queue tests must not access a network or GPU")
monkeypatch.setattr(aiohttp, "ClientSession", forbidden)
monkeypatch.setattr(socket, "create_connection", forbidden)
monkeypatch.setattr(socket.socket, "connect", forbidden)
monkeypatch.setattr(socket.socket, "connect_ex", forbidden)
monkeypatch.setattr(torch.cuda, "init", forbidden)
monkeypatch.setattr(torch.cuda, "current_device", forbidden)
h = types.SimpleNamespace(q=q, ps=ps, pool=pool, config=config, manager=manager,
client=client, compiler=compiler, packer=packer,
validator=validator, handler=handler)
yield h
@pytest.fixture
def media_paths(tmp_path, monkeypatch):
paths = types.SimpleNamespace(get_annotated_filepath=lambda name: str(tmp_path / name))
for kind in ("input", "output", "temp"):
setattr(paths, "get_" + kind + "_directory", lambda: str(tmp_path))
monkeypatch.setitem(sys.modules, "folder_paths", paths)
return tmp_path
def run(h, prompt=None, **kwargs):
return asyncio.run(h.q.orchestrate_distributed_execution(
prompt if prompt is not None else text_prompt(), None, "client", **kwargs))
def assert_no_dispatch(h, *, no_probe=False):
h.q.queue_prompt_payload.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
h.manager.start.assert_not_awaited()
h.manager.cancel.assert_not_awaited()
assert h.ps.distributed_pending_jobs == {}
if no_probe:
h.q.select_active_workers.assert_not_awaited()
h.q.select_least_busy_worker.assert_not_awaited()
h.q.fetch_worker_has_node.assert_not_awaited()
h.q.sync_worker_media.assert_not_awaited()
@pytest.mark.parametrize("delegate", [True, False])
def test_real_compile_pack_serialize_and_canonical_bindings(harness, media_paths, delegate):
h = harness
image_path = media_paths / "picture.png"
Image.new("RGBA", (2, 3), color=(17, 29, 41, 53)).save(image_path)
image_bytes = image_path.read_bytes()
prompt = text_prompt()
prompt["image"] = node("LoadImage", image="picture.png")
prompt["collector"]["inputs"].update(images=["image", 0], masks=["image", 1])
prompt["collector2"] = node("DistributedMediaCollector", texts=["text", 0])
prompt["unused"] = node("PreviewImage", images=["image", 0])
source = copy.deepcopy(prompt)
with patch.object(h.compiler, "compile_media_prompt", wraps=h.compiler.compile_media_prompt) as compile_call, \
patch.object(h.packer, "prepare_media_inputs", wraps=h.packer.prepare_media_inputs) as pack_call, \
patch.object(h.client, "serialize_request", wraps=h.client.serialize_request) as serialize_call, \
patch.object(h.q, "prune_prompt_for_worker", side_effect=AssertionError("typed graph must bypass legacy pruning")):
result = run(h, prompt, delegate_master=delegate)
assert result == ("master-prompt", 7, 2, {})
assert compile_call.call_count == pack_call.await_count == 2
assert serialize_call.call_count >= 2
assert prompt == source
h.q.dispatch_worker_prompt.assert_not_awaited()
jobs = h.manager.start.call_args.args[2]
assert [j["worker"]["id"] for j in jobs] == ["serverless:pool:1", "serverless:pool:2"]
assert len({j["media"]["request_id"] for j in jobs}) == 2
master = h.q.queue_prompt_payload.call_args.args[0]
for job in jobs:
media = job["media"]
assert str(uuid.UUID(media["request_id"])) == media["request_id"]
assert uuid.UUID(media["request_id"]).version == 4
assert media["profile"] == PROFILE and media["version"] == 1
assert media["outputs"] == job["manifest"]
assert media["outputs"] is not job["manifest"]
assert media["artifacts"] == media["inputs"] == []
assert {b["socket"] for b in job["manifest"]} == {"texts", "images", "masks"}
assert len(job["manifest"]) == 4
assert base64.b64decode(job["images"][0]["image"]) == image_bytes
assert job["workflow"]["image"]["inputs"]["image"] == job["images"][0]["name"]
assert not ({"save", "unused", "collector", "collector2"} & set(job["workflow"]))
body = h.client.serialize_request(job["workflow"], job["images"], 30, media=media)
wire = json.loads(body)
assert wire["input"] == {"workflow": job["workflow"], "images": job["images"], "media": media}
assert wire["policy"] == {"executionTimeout": 30000, "ttl": 30000}
_, _, _, targets = h.handler._validate_request(wire["input"])
assert targets == [b["node_id"] for b in job["manifest"]]
assert h.validator.validate_prepared_job(job) is True
for binding in job["manifest"]:
assert binding["participant_id"] == job["worker"]["id"]
assert binding["job_id"] == master[binding["collector_id"]]["inputs"]["multi_job_id"]
controls = job["workflow"][binding["node_id"]]["inputs"]
assert json.loads(controls["binding_json"]) == binding
assert controls["request_id"] == media["request_id"]
assert master["collector"]["inputs"]["delegate_only"] is delegate
assert json.loads(master["collector"]["inputs"]["expected_media"]) == ["images", "masks", "texts"]
if delegate:
assert "text" not in master and "image" not in master
assert not any(name in master["collector"]["inputs"] for name in ("images", "audio", "masks", "texts", "video"))
assert not any(n["class_type"] == "DistributedEmptyImage" for n in master.values())
else:
assert master["collector"]["inputs"]["texts"] == ["text", 0]
assert "text" in master and "image" in master
def test_delegate_disconnects_every_media_socket(harness):
prompt = text_prompt()
for index, name in enumerate(("images", "audio", "masks", "video")):
source_id = "source-" + name
prompt[source_id] = node("GeneratedMedia", seed=index)
prompt["collector"]["inputs"][name] = [source_id, 0]
original = copy.deepcopy(prompt)
run(harness, prompt)
master = harness.q.queue_prompt_payload.call_args.args[0]
inputs = master["collector"]["inputs"]
sockets = {"images", "audio", "masks", "texts", "video"}
assert set(json.loads(inputs["expected_media"])) == sockets
assert not sockets.intersection(inputs)
assert set(master) == {"collector", "save"}
for job in harness.manager.start.call_args.args[2]:
assert {b["socket"] for b in job["manifest"]} == sockets
assert {"text", "source-images", "source-audio", "source-masks", "source-video"} <= set(job["workflow"])
assert prompt == original
@pytest.mark.parametrize("selected", ["master", "worker"])
def test_typed_load_balance_selects_master_or_one_serverless_slot(harness, selected):
h = harness
prompt = text_prompt()
prompt["collector"]["inputs"]["load_balance"] = "true"
h.q.select_least_busy_worker.side_effect = lambda candidates, **kw: next(
w for w in candidates if (w["id"] == "master") == (selected == "master"))
result = run(h, prompt, delegate_master=False)
h.q.select_least_busy_worker.assert_awaited_once()
master = h.q.queue_prompt_payload.call_args.args[0]
if selected == "master":
assert result[2] == 0
h.manager.start.assert_not_awaited()
assert master["collector"]["inputs"]["delegate_only"] is False
assert "text" in master
else:
assert result[2] == 1
assert len(h.manager.start.call_args.args[2]) == 1
assert master["collector"]["inputs"]["delegate_only"] is True
assert "text" not in master
@pytest.mark.parametrize("profile", [None, "legacy-image-v1", PROFILE])
def test_explicit_profile_compatibility_rejected_before_probe(harness, profile):
h = harness
prompt = legacy_prompt() if profile == PROFILE else text_prompt()
if profile is None:
h.pool.pop("media_profile")
else:
h.pool["media_profile"] = profile
with pytest.raises(ValueError, match="requires|require"):
run(h, prompt)
assert_no_dispatch(h, no_probe=True)
@pytest.mark.parametrize("profile", [None, "auto", "comfy-media-v2", False])
def test_invalid_explicit_profiles_fail_closed(harness, profile):
harness.pool["media_profile"] = profile
with pytest.raises(ValueError, match="media_profile"):
run(harness)
assert_no_dispatch(harness, no_probe=True)
@pytest.mark.parametrize("direct_type", ["local", "remote", "cloud"])
@pytest.mark.parametrize("explicit", [False, True])
def test_direct_workers_never_silently_skipped_or_fallback(harness, direct_type, explicit):
h = harness
h.config["workers"].append({"id": "direct", "type": direct_type, "host": "offline.invalid",
"enabled": not explicit, "port": 8188})
prompt = text_prompt()
prompt["collector"]["inputs"]["load_balance"] = True
kwargs = {"enabled_worker_ids": ["pool", "direct"]} if explicit else {}
with pytest.raises(ValueError, match="direct workers"):
run(h, prompt, delegate_master=False, **kwargs)
assert_no_dispatch(h, no_probe=True)
def test_disabled_or_unselected_incompatible_workers_do_not_block(harness):
h = harness
h.config["workers"] += [
{"id": "direct", "type": "local", "enabled": False},
dict(h.pool, id="legacy", media_profile="legacy-image-v1", enabled=False),
]
assert run(h)[2] == 2
h.pool["enabled"] = False
assert run(h, enabled_worker_ids=["pool"])[2] == 2
@pytest.mark.parametrize("other", ["DistributedCollector", "UltimateSDUpscaleDistributed", "DistributedImageListStart"])
@pytest.mark.parametrize("workers", [True, False])
def test_mixed_collectors_or_image_list_rejected_even_master_only(harness, other, workers):
h = harness
if not workers:
h.config["workers"] = []
prompt = text_prompt()
prompt["incompatible"] = node(other)
with pytest.raises(ValueError, match="cannot be mixed"):
run(h, prompt, delegate_master=False)
assert_no_dispatch(h, no_probe=True)
def test_mixed_pool_profiles_rejected_before_load_balancing(harness):
h = harness
h.config["workers"].append(dict(h.pool, id="legacy", media_profile="legacy-image-v1"))
prompt = text_prompt()
prompt["collector"]["inputs"]["load_balance"] = True
with pytest.raises(ValueError, match="requires"):
run(h, prompt, delegate_master=False)
assert_no_dispatch(h, no_probe=True)
def test_typed_master_only_is_supported_without_serverless_work(harness):
h = harness
h.config["workers"] = []
assert run(h, delegate_master=False)[2] == 0
master = h.q.queue_prompt_payload.call_args.args[0]
assert master["collector"]["inputs"]["enabled_worker_ids"] == "[]"
assert master["collector"]["inputs"]["texts"] == ["text", 0]
h.manager.start.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
@pytest.mark.parametrize("offline", [False, True])
def test_typed_delegate_requires_available_worker_without_master_fallback(harness, offline):
h = harness
if offline:
h.q.select_active_workers.side_effect = lambda *a, **kw: ([], False)
else:
h.config["workers"] = []
with pytest.raises(RuntimeError, match="comfy-media-v1"):
run(h)
assert_no_dispatch(h)
def test_unknown_explicit_typed_worker_is_rejected(harness):
with pytest.raises(ValueError, match="Unknown explicitly selected"):
run(harness, enabled_worker_ids=["unknown"], delegate_master=False)
assert_no_dispatch(harness, no_probe=True)
@pytest.mark.parametrize("bad", [" pool ", "pool white", "pool\n"])
def test_pool_media_identity_is_not_repaired(harness, bad):
harness.pool["id"] = bad
with pytest.raises(ValueError, match="identifier"):
run(harness)
assert_no_dispatch(harness, no_probe=True)
def test_integer_source_node_identity_is_not_normalized(harness):
prompt = text_prompt()
prompt[1] = prompt.pop("text")
prompt["collector"]["inputs"]["texts"] = [1, 0]
with pytest.raises(ValueError, match="identifier"):
run(harness, prompt)
assert_no_dispatch(harness, no_probe=True)
@pytest.mark.parametrize("broken", ["dangling-unused", "literal", "missing-file", "bad-loader-controls", "oversize"])
def test_compilation_packing_and_size_errors_precede_all_dispatch(harness, media_paths, broken):
h = harness
prompt = text_prompt()
if broken == "dangling-unused":
prompt["unused"] = node("PreviewImage", images=["missing", 0])
elif broken == "literal":
prompt["collector"]["inputs"]["texts"] = "must use a linked string node"
elif broken in ("missing-file", "bad-loader-controls"):
prompt["audio"] = node("LoadAudio", audio="missing.wav")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
if broken == "bad-loader-controls":
prompt["audio"]["inputs"]["duration"] = 5
else:
prompt["text"]["inputs"]["value"] = "雪" * 1000
original = copy.deepcopy(prompt)
with patch.object(h.client, "MAX_REQUEST_BYTES", 100 if broken == "oversize" else h.client.MAX_REQUEST_BYTES):
with pytest.raises((ValueError, OSError)):
run(h, prompt)
assert prompt == original
assert_no_dispatch(h)
@pytest.mark.parametrize("corruption", ["manifest", "participant", "request-control", "binding-control", "profile", "input-ref"])
def test_whole_batch_identity_preflight_before_master_and_manager(harness, corruption):
h = harness
real_prepare = h.q._prepare_worker_payload
async def corrupt_later(*args, **kwargs):
worker, prepared = await real_prepare(*args, **kwargs)
if worker["id"].endswith(":2"):
binding = prepared["manifest"][0]
if corruption == "manifest":
binding["job_id"] = "wrong-job"
elif corruption == "participant":
prepared["worker"] = dict(worker, id="different-participant")
elif corruption == "request-control":
prepared["workflow"][binding["node_id"]]["inputs"]["request_id"] = "different-request"
elif corruption == "binding-control":
prepared["workflow"][binding["node_id"]]["inputs"]["binding_json"] = "{}"
elif corruption == "profile":
prepared["worker"] = dict(worker, media_profile="legacy-image-v1")
else:
prepared["media"]["inputs"] = [{"node_id": "missing", "input_name": "texts", "artifact_id": "unknown"}]
return worker, prepared
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
with pytest.raises(ValueError):
run(h)
h.manager.ensure_capacity.assert_not_called()
assert_no_dispatch(h)
def test_batch_preflight_also_precedes_regular_direct_dispatch(harness):
h = harness
h.pool.pop("media_profile")
h.config["workers"].append({"id": "direct", "type": "local", "enabled": True})
real_prepare = h.q._prepare_worker_payload
async def corrupt_later(*args, **kwargs):
worker, prepared = await real_prepare(*args, **kwargs)
if worker["id"] == "serverless:pool:2":
prepared["worker"] = dict(worker, media_profile=PROFILE)
return worker, prepared
with patch.object(h.q, "_prepare_worker_payload", side_effect=corrupt_later):
with pytest.raises(ValueError, match="mismatch"):
run(h, legacy_prompt())
assert_no_dispatch(h)
def test_capacity_preflight_precedes_collector_state_and_master(harness):
h = harness
h.manager.ensure_capacity.side_effect = ValueError("capacity exhausted")
with pytest.raises(ValueError, match="capacity"):
run(h)
assert_no_dispatch(h)
def test_master_queue_failure_cleans_queues_and_never_starts_manager(harness):
h = harness
h.q.queue_prompt_payload.side_effect = ValueError("invalid master graph")
with pytest.raises(ValueError, match="invalid master"):
run(h)
assert h.ps.distributed_pending_jobs == {}
h.manager.start.assert_not_awaited()
h.q.dispatch_worker_prompt.assert_not_awaited()
def test_manager_setup_failure_emits_one_terminal_per_collector_participant(harness):
h = harness
prompt = text_prompt()
prompt["image"] = node("EmptyImage", width=2, height=2, batch_size=1, color=0)
prompt["collector"]["inputs"]["images"] = ["image", 0]
h.manager.start.side_effect = ValueError("journal failure")
with pytest.raises(ValueError, match="journal failure"):
run(h, prompt)
h.manager.cancel.assert_awaited_once()
queued = next(iter(h.ps.distributed_pending_jobs.values()))
assert queued.qsize() == 2 # two participants, not four per-socket completions
messages = [queued.get_nowait(), queued.get_nowait()]
assert {m["worker_id"] for m in messages} == {"serverless:pool:1", "serverless:pool:2"}
assert all(m["is_last"] is True and m["serverless_error"] for m in messages)
@pytest.mark.parametrize("profile", [None, "legacy-image-v1"])
def test_legacy_envelope_and_compiler_are_unchanged(harness, profile):
h = harness
if profile is None:
h.pool.pop("media_profile")
else:
h.pool["media_profile"] = profile
prompt = legacy_prompt()
original = copy.deepcopy(prompt)
with patch.object(h.compiler, "compile_media_prompt", side_effect=AssertionError("legacy must not use typed compiler")), \
patch.object(h.packer, "prepare_media_inputs", side_effect=AssertionError("legacy must not use typed packer")):
assert run(h, prompt)[2] == 2
for job in h.manager.start.call_args.args[2]:
assert "media" not in job
assert any(n["class_type"] == "SaveImage" for n in job["workflow"].values())
expected = {"input": {"workflow": job["workflow"], "images": job["images"]},
"policy": {"executionTimeout": 30000, "ttl": 30000}}
assert h.client.serialize_request(job["workflow"], job["images"], 30) == json.dumps(
expected, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode()
assert prompt == original
def test_legacy_no_worker_delegate_fallback_is_preserved(harness):
h = harness
h.config["workers"] = []
assert run(h, legacy_prompt())[2] == 0
assert h.q.queue_prompt_payload.call_args.args[0]["collector"]["inputs"]["delegate_only"] is False
def test_native_audio_compile_pack_serialized_input_cpu(harness, media_paths):
"""Exercise real PyAV + core PCM conversion, not a mocked packing result."""
if os.environ.get("_MEDIA_QUEUE_CORE_TEST") != "1":
core = os.environ.get("COMFYUI_SOURCE_ROOT")
if not core:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native audio queue integration")
env = dict(os.environ)
env.pop("PYTHONPATH", None)
env.pop("COMFYUI_SOURCE_ROOT", None)
scratch = Path.home() / ".hermes/cache/scratch"
scratch.mkdir(parents=True, exist_ok=True)
env.update(_MEDIA_QUEUE_CORE_TEST="1", CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
PYTHONDONTWRITEBYTECODE="1", TMPDIR=str(scratch))
bootstrap = """
import sys
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
import torch
def forbidden(*args, **kwargs):
raise AssertionError('CPU audio integration may not initialize a GPU')
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
torch.cuda.init = forbidden
torch.cuda.current_device = forbidden
torch.cuda._lazy_init = forbidden
import pytest
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider', '-k',
'test_native_audio_compile_pack_serialized_input_cpu']))
"""
result = subprocess.run([sys.executable, "-c", bootstrap, core, str(Path(__file__).resolve())],
env=env, capture_output=True, text=True, timeout=120)
assert result.returncode == 0, result.stdout + result.stderr
assert "1 passed" in result.stdout
return
h = harness
h.pool["jobs_per_run"] = 1
path = media_paths / "tone.wav"
pcm = torch.tensor([-100, 0, 100, 32767, -32768, 7], dtype=torch.int16)
with wave.open(str(path), "wb") as stream:
stream.setnchannels(1)
stream.setsampwidth(2)
stream.setframerate(8000)
stream.writeframes(pcm.numpy().tobytes())
original_bytes = path.read_bytes()
prompt = text_prompt()
prompt["audio"] = node("LoadAudio", audio="tone.wav")
prompt["collector"]["inputs"]["audio"] = ["audio", 0]
original = copy.deepcopy(prompt)
run(h, prompt)
job = h.manager.start.call_args.args[2][0]
wire = json.loads(h.client.serialize_request(job["workflow"], job["images"], 30, media=job["media"]))
artifacts = wire["input"]["media"]["artifacts"]
refs = wire["input"]["media"]["inputs"]
assert len(artifacts) == len(refs) == 1
binding = next(b for b in job["manifest"] if b["socket"] == "audio")
assert refs == [{"node_id": binding["node_id"], "input_name": "audio", "artifact_id": artifacts[0]["id"]}]
contract = importlib.import_module(PKG + ".serverless_bridge.contract")
codec = importlib.import_module(PKG + ".serverless_bridge.codec")
payload = asyncio.run(contract.resolve_artifact(artifacts[0]))
value = codec.decode_payload("AUDIO", artifacts[0]["descriptor"], payload)
assert value["sample_rate"] == 8000
assert value["waveform"].shape == (1, 1, 6)
assert torch.equal(value["waveform"], (pcm.to(torch.float32) / 32768).reshape(1, 1, -1))
assert path.read_bytes() == original_bytes
assert prompt == original
assert h.validator.validate_prepared_job(job) is True
+76
View File
@@ -0,0 +1,76 @@
"""Collector controls -> compiler -> endpoint validation, no network or GPU."""
import copy
import json
import pytest
from test_serverless_media_prompt import compiler, _prompt
from test_media_handler import handler
from test_media_collector import module, collect, native_media
FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1')
def compile_request(fmt='exact', quality=90):
prompt = _prompt()
prompt['source:1'] = {'class_type': 'Source', 'inputs': {}}
prompt['collector:1']['inputs'].update(transfer_format=fmt, transfer_quality=quality,
masks=['source:1', 1], texts=['source:1', 2])
original = copy.deepcopy(prompt)
workflow, outputs = compiler.compile_media_prompt(prompt, 'request:1')
assert prompt == original
return {'workflow': workflow, 'images': [], 'media': {
'version': 1, 'profile': 'comfy-media-v1', 'request_id': 'request:1',
'artifacts': [], 'inputs': [], 'outputs': outputs}}
@pytest.mark.parametrize('fmt', FORMATS)
def test_compiler_and_handler_preserve_image_policy_only(fmt):
data = compile_request(fmt, 73)
for binding in data['media']['outputs']:
inputs = data['workflow'][binding['node_id']]['inputs']
if binding['kind'] == 'IMAGE' and fmt != 'exact':
assert inputs['transfer_format'] == fmt
assert inputs['transfer_quality'] == 73
else:
assert 'transfer_format' not in inputs and 'transfer_quality' not in inputs
handler._validate_request(data)
@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('',90),('jpeg',True),('jpeg',0),
('jpeg',101),('jpeg','90'),(['source:1',0],90)])
def test_compiler_rejects_invalid_or_linked_controls(fmt, quality):
with pytest.raises(ValueError):
compile_request(fmt, quality)
@pytest.mark.parametrize('field,value', [('transfer_format','unknown'),('transfer_quality',0),
('transfer_quality',True),('transfer_quality','90')])
def test_handler_rejects_tampered_policy(field, value):
data = compile_request('jpeg')
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'IMAGE')
data['workflow'][binding['node_id']]['inputs'][field] = value
with pytest.raises(ValueError):
handler._validate_request(data)
def test_handler_rejects_image_controls_on_mask_export():
data = compile_request()
binding = next(b for b in data['media']['outputs'] if b['kind'] == 'MASK')
data['workflow'][binding['node_id']]['inputs']['transfer_format'] = 'jpeg'
with pytest.raises(ValueError):
handler._validate_request(data)
@pytest.mark.parametrize('fmt', FORMATS)
def test_master_passthrough_stays_native_for_every_format(module, fmt):
media = native_media()
actual = collect(module, **media, transfer_format=[fmt], transfer_quality=[73])
for values, original in zip(actual, media.values()):
assert values is original
@pytest.mark.parametrize('quality', [True, 0, 101, '90', [90,91]])
def test_node_rejects_bad_quality_even_without_remote_participants(module, quality):
with pytest.raises(ValueError):
collect(module, texts=['hello'], transfer_quality=quality)
+562
View File
@@ -0,0 +1,562 @@
"""Native ComfyUI VIDEO tests, isolated from the custom-node import graph.
The outer pytest process starts one bounded CPU-only subprocess. Inside it the
real pinned core is imported *after* cli_args.args.cpu=True; no fake VideoInput
classes, media transcodes, providers, server, or GPU queries are needed.
"""
import os
from pathlib import Path
import subprocess
import sys
import pytest
if os.environ.get("_DISTRIBUTED_VIDEO_CODEC_CORE_TEST") != "1":
def test_native_video_codec_in_cpu_subprocess():
configured = os.environ.get("COMFYUI_SOURCE_ROOT", os.environ.get("COMFYUI_VIDEO_TEST_ROOT"))
if configured is None:
pytest.skip("Set COMFYUI_SOURCE_ROOT for CPU native VIDEO acceptance")
core = Path(configured).expanduser().resolve()
assert configured and (core / "comfy_api/latest/_input_impl/video_types.py").is_file(), (
"Configured COMFYUI_SOURCE_ROOT (or COMFYUI_VIDEO_TEST_ROOT) must provide native video classes"
)
env = dict(os.environ)
scratch = Path(os.environ.get("TMPDIR", str(Path.home() / ".hermes/cache/scratch")))
scratch.mkdir(parents=True, exist_ok=True)
env["TMPDIR"] = str(scratch)
for name in ("PYTHONPATH", "COMFYUI_SOURCE_ROOT"):
env.pop(name, None)
env.update(CUDA_VISIBLE_DEVICES="", HIP_VISIBLE_DEVICES="",
_DISTRIBUTED_VIDEO_CODEC_CORE_TEST="1", PYTHONDONTWRITEBYTECODE="1")
bootstrap = """
import sys
sys.path.insert(0, sys.argv[1])
import comfy.cli_args
comfy.cli_args.args.cpu = True
import torch
def forbidden(*args, **kwargs):
raise AssertionError('VIDEO tests must not query or initialize a GPU')
# Third-party core imports probe availability even in CPU mode: answer without
# executing a GPU query, then forbid even these probes during codec execution.
torch.cuda.is_available = lambda: False
torch.cuda.device_count = lambda: 0
for name in ('init', '_lazy_init', 'current_device'):
setattr(torch.cuda, name, forbidden)
import comfy_api.latest._input_impl.video_types
for name in ('is_available', 'device_count'):
setattr(torch.cuda, name, forbidden)
import pytest
raise SystemExit(pytest.main([sys.argv[2], '-q', '-p', 'no:cacheprovider']))
"""
result = subprocess.run(
[sys.executable, "-c", bootstrap, str(core), str(Path(__file__).resolve())],
env=env, text=True, capture_output=True, timeout=150,
cwd=Path(__file__).resolve().parents[1],
)
print(result.stdout)
assert result.returncode == 0, result.stdout + result.stderr
else:
import copy
from fractions import Fraction
import importlib
import io
import json
import math
import types
import av
import numpy as np
import torch
from comfy_api.latest._input import VideoInput
from comfy_api.latest._input_impl.video_types import (
VideoFromComponents, VideoFromFile, VideoFromList,
)
from comfy_api.latest._util import VideoComponents
# Only this sibling package is synthetic; all native media types are real.
package_name = "_distributed_video_codec_test"
package = types.ModuleType(package_name)
package.__path__ = [str(Path(__file__).resolve().parents[1] / "serverless_bridge")]
sys.modules[package_name] = package
codec = importlib.import_module(package_name + ".video_codec")
def audio(dtype=torch.float64):
return {"waveform": torch.tensor([[[0.123456789, -0.0, -1.5, 2.0]]], dtype=dtype, device="cpu"),
"sample_rate": 48000}
def components(dtype=torch.float64, *, metadata=None, alpha=True, sound=True):
image = torch.arange(36, device="cpu", dtype=torch.float64).reshape(2, 2, 3, 3) / 997
image[0, 0, 0, 0] = -0.0
return VideoFromComponents(VideoComponents(
images=image.to(dtype), frame_rate=Fraction(30000, 1001),
audio=audio(dtype) if sound else None,
alpha=torch.full((2, 2, 3), 0.123456789, dtype=dtype, device="cpu") if alpha else None,
metadata=metadata,
), bit_depth=10, color_space="HDR")
def exact_tensor(actual, expected):
assert actual.device.type == "cpu"
assert actual.dtype == expected.dtype
assert actual.shape == expected.shape
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
def file_value(data=b"opaque original media bytes", **kwargs):
return VideoFromFile(io.BytesIO(data), **kwargs)
def forbid(*args, **kwargs):
raise AssertionError("Unexpected materialization/transcoding/allocation")
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
@pytest.mark.parametrize("alpha,sound", [(False, False), (True, False), (False, True), (True, True)])
def test_components_exact_roundtrip(dtype, alpha, sound, monkeypatch):
metadata = {"title": "雪 🎞", "flags": [True, None, -4, 0.125], "nested": {"empty": []}}
source = components(dtype, metadata=metadata, alpha=alpha, sound=sound)
expected = source.get_components()
monkeypatch.setattr(VideoFromComponents, "get_components", forbid)
monkeypatch.setattr(VideoFromComponents, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert descriptor["encoding"] == "native-video-v1"
assert descriptor["representation"] == "components"
assert descriptor["frame_rate"] == {"numerator": 30000, "denominator": 1001}
# This is the actual JSON wire descriptor, not a Python-only shortcut.
descriptor = json.loads(json.dumps(descriptor, allow_nan=False))
result = codec.decode_video(descriptor, payload)
assert type(result) is VideoFromComponents
assert isinstance(result, VideoInput)
actual = vars(result)["_VideoFromComponents__components"]
exact_tensor(actual.images, expected.images)
assert actual.frame_rate == expected.frame_rate
assert type(actual.frame_rate) is Fraction
assert actual.metadata == metadata
assert actual.metadata is not metadata
assert vars(result)["_VideoFromComponents__bit_depth"] == 10
assert vars(result)["_VideoFromComponents__color_space"] == "HDR"
if alpha:
exact_tensor(actual.alpha, expected.alpha)
else:
assert actual.alpha is None
if sound:
exact_tensor(actual.audio["waveform"], expected.audio["waveform"])
assert actual.audio["sample_rate"] == expected.audio["sample_rate"]
else:
assert actual.audio is None
actual.images.fill_(99)
assert not torch.equal(actual.images, expected.images)
actual.metadata["nested"]["empty"].append("new")
assert descriptor["metadata"]["nested"]["empty"] == []
assert metadata["nested"]["empty"] == []
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
def test_native_bhw1_alpha_shape_is_preserved(dtype):
source = components(dtype)
original = source.get_components()
vars(source)["_VideoFromComponents__components"].alpha = original.alpha.unsqueeze(-1)
descriptor, payload = codec.encode_video(source)
result = codec.decode_video(descriptor, payload).get_components()
exact_tensor(result.alpha, original.alpha.unsqueeze(-1))
assert descriptor["alpha"]["descriptor"]["shape"] == [2, 2, 3, 1]
@pytest.mark.parametrize("start,duration,crop", [
(-1.25, -2.5, (-3, 5, 7, -9)), (0, 0, None), (-0.0, 0.0, (0, 0, 2, 2)),
])
def test_file_preserves_bytes_cursor_signed_trim_crop_and_ownership(start, duration, crop, monkeypatch):
source = file_value(start_time=start, duration=duration, crop=crop)
stream = vars(source)["_VideoFromFile__file"]
stream.seek(7)
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert stream.tell() == 7
assert payload == stream.getvalue()
assert descriptor["start_time"] == start
assert descriptor["duration"] == duration
result = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
second = codec.decode_video(descriptor, payload)
attrs = vars(result)
assert type(result) is VideoFromFile
assert attrs["_VideoFromFile__start_time"] == start
assert type(attrs["_VideoFromFile__start_time"]) is type(start)
assert attrs["_VideoFromFile__duration"] == duration
assert attrs["_VideoFromFile__crop"] == crop
if type(start) is float:
assert math.copysign(1, attrs["_VideoFromFile__start_time"]) == math.copysign(1, start)
decoded = attrs["_VideoFromFile__file"]
assert type(decoded) is io.BytesIO
assert decoded is not stream
stream.close()
assert decoded.getvalue() == payload
decoded.write(b"changed")
assert vars(second)["_VideoFromFile__file"].getvalue() == payload
def test_real_tiny_pyav_video_is_still_readable_after_source_closed():
buffer = io.BytesIO()
with av.open(buffer, "w", format="matroska") as container:
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
# Width 32 avoids a native fillborders filter limitation on tiny
# unaligned YUV frames; this is still a tiny CPU-only fixture.
stream.width = stream.height = 32
stream.pix_fmt = "yuv420p"
for value in (16, 64, 128):
frame = av.VideoFrame.from_ndarray(np.full((32, 32, 3), value, dtype=np.uint8), format="rgb24")
for packet in stream.encode(frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
data = buffer.getvalue()
source = VideoFromFile(buffer, crop=(0, 0, 2, 2))
descriptor, payload = codec.encode_video(source)
expected = source.get_components()
result = codec.decode_video(descriptor, payload)
buffer.close()
assert payload == data
actual = result.get_components()
exact_tensor(actual.images, expected.images)
assert actual.frame_rate == expected.frame_rate
assert tuple(actual.images.shape) == (3, 2, 2, 3)
def test_segments_order_complete_audio_no_transcode_and_independent_buffers(monkeypatch):
values = [file_value(b"first", start_time=-2), file_value(b"second", duration=-4, crop=(1, 2, 3, 4))]
source = VideoFromList([VideoFromList(values[:1]), values[1]], complete_audio=audio())
monkeypatch.setattr(VideoFromFile, "get_components", forbid)
monkeypatch.setattr(VideoFromFile, "save_to", forbid)
monkeypatch.setattr(VideoFromList, "get_components", forbid)
monkeypatch.setattr(VideoFromList, "save_to", forbid)
descriptor, payload = codec.encode_video(source)
assert descriptor["representation"] == "segments"
assert payload.startswith(b"firstsecond")
assert len(descriptor["segments"]) == 2
decoded = codec.decode_video(json.loads(json.dumps(descriptor)), payload)
assert type(decoded) is VideoFromList
assert [vars(v)["_VideoFromFile__file"].getvalue() for v in decoded.videos] == [b"first", b"second"]
exact_tensor(decoded.complete_audio["waveform"], source.complete_audio["waveform"])
for old, new in zip(source.videos, decoded.videos):
assert old is not new
vars(old)["_VideoFromFile__file"].close()
assert not vars(new)["_VideoFromFile__file"].closed
assert vars(decoded.videos[0])["_VideoFromFile__start_time"] == -2
assert vars(decoded.videos[1])["_VideoFromFile__crop"] == (1, 2, 3, 4)
def test_path_roots_symlinks_and_no_path_in_wire(tmp_path, monkeypatch):
allowed = tmp_path / "allowed"
allowed.mkdir()
media = allowed / "private-name.mkv"
media.write_bytes(b"media bytes")
sibling = tmp_path / "allowed-sibling"
sibling.mkdir()
secret = sibling / "secret.mkv"
secret.write_bytes(b"private")
(allowed / "escape.mkv").symlink_to(secret)
(allowed / "inside.mkv").symlink_to(media)
descriptor, payload = codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
assert payload == b"media bytes"
assert "private-name" not in json.dumps(descriptor)
assert codec.encode_video(VideoFromFile(str(allowed / "inside.mkv")), allowed_roots=[allowed])[1] == payload
for path in (secret, allowed / "escape.mkv", allowed, allowed / "missing"):
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(path)), allowed_roots=[allowed])
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[])
import folder_paths
for name in ("get_input_directory", "get_output_directory", "get_temp_directory"):
monkeypatch.setattr(folder_paths, name, lambda: str(allowed))
assert codec.encode_video(VideoFromFile(str(media)))[1] == payload
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(secret)))
result = codec.decode_video(descriptor, payload)
media.unlink()
assert vars(result)["_VideoFromFile__file"].getvalue() == b"media bytes"
def test_path_special_files_and_symlink_swap_are_rejected(tmp_path, monkeypatch):
fifo = tmp_path / "fifo"
os.mkfifo(fifo)
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(fifo)), allowed_roots=[tmp_path])
allowed = tmp_path / "allowed"
allowed.mkdir()
actual = allowed / "subdir"
actual.mkdir()
media = actual / "clip.mkv"
media.write_bytes(b"allowed")
outside = tmp_path / "outside"
outside.mkdir()
(outside / "clip.mkv").write_bytes(b"forbidden")
real_resolve = codec.Path.resolve
swapped = False
def swap_after_resolution(path, *args, **kwargs):
nonlocal swapped
resolved = real_resolve(path, *args, **kwargs)
if path == media and not swapped:
swapped = True
actual.rename(allowed / "old")
actual.symlink_to(outside, target_is_directory=True)
return resolved
monkeypatch.setattr(codec.Path, "resolve", swap_after_resolution)
with pytest.raises(ValueError):
codec.encode_video(VideoFromFile(str(media)), allowed_roots=[allowed])
assert swapped
def test_file_size_budget_checked_before_read(tmp_path, monkeypatch):
path = tmp_path / "large.mkv"
path.write_bytes(b"12345")
monkeypatch.setattr(codec.os, "read", forbid)
with pytest.raises(ValueError, match="limit|budget"):
codec.encode_video(VideoFromFile(str(path)), max_bytes=4, allowed_roots=[tmp_path])
@pytest.mark.parametrize("limit", [True, False, -1, 1.5, None])
def test_invalid_limits(limit):
with pytest.raises(ValueError):
codec.encode_video(file_value(), max_bytes=limit)
with pytest.raises(ValueError):
codec.decode_video({}, b"", max_bytes=limit)
def test_aggregate_budget_and_exact_boundary():
value = VideoFromList([file_value(b"12"), file_value(b"345")])
descriptor, payload = codec.encode_video(value, max_bytes=5)
assert payload == b"12345"
codec.decode_video(descriptor, payload, max_bytes=5)
with pytest.raises(ValueError):
codec.encode_video(value, max_bytes=4)
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload, max_bytes=4)
stream = vars(value.videos[0])["_VideoFromFile__file"]
stream.seek(1)
with pytest.raises(ValueError):
codec.encode_video(value.videos[0], max_bytes=1)
assert stream.tell() == 1
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
def test_unknown_subclasses_never_materialized(kind):
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
value = object.__new__(cls)
with pytest.raises(ValueError, match="class|type|unsupported|Unsupported"):
codec.encode_video(value)
with pytest.raises(ValueError):
codec.encode_video(object())
@pytest.mark.parametrize("kind", ["components", "file", "segments"])
def test_bridge_rejects_subclasses_without_materializing(kind, monkeypatch):
nodes = importlib.import_module(package_name + ".nodes")
bases = {"components": VideoFromComponents, "file": VideoFromFile, "segments": VideoFromList}
cls = type("UnknownVideo", (bases[kind],), {"get_components": forbid, "save_to": forbid})
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
with pytest.raises(ValueError, match="Unsupported native VIDEO class"):
nodes.encode_value("VIDEO", object.__new__(cls))
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
@pytest.mark.parametrize("kind", ["components", "segments", "file"])
def test_bridge_stages_cpu_tensors_without_materialization(dtype, kind, monkeypatch):
nodes = importlib.import_module(package_name + ".nodes")
if kind == "components":
source = components(dtype, metadata={"nested": ["雪", None, True]})
state = vars(source)["_VideoFromComponents__components"]
state.alpha = state.alpha.unsqueeze(-1)
tensors = [state.images, state.alpha, state.audio["waveform"]]
elif kind == "segments":
source = VideoFromList([file_value(start_time=-1.5, duration=-2, crop=(-1, 2, 3, 4))],
complete_audio=audio(dtype))
tensors = [source.complete_audio["waveform"]]
else:
source = file_value(start_time=-0.0, duration=-2, crop=(-1, 2, 3, 4))
tensors = []
for tensor in tensors:
tensor.requires_grad_(True)
expected = codec.encode_video(source)
detached, staged = [], []
real_detach, real_cpu = torch.Tensor.detach, torch.Tensor.cpu
def detach(tensor):
detached.append(id(tensor))
return real_detach(tensor)
def cpu(tensor, *args, **kwargs):
assert not tensor.requires_grad
staged.append(tensor)
return real_cpu(tensor, *args, **kwargs)
for cls in (VideoFromComponents, VideoFromFile, VideoFromList):
monkeypatch.setattr(cls, "get_components", forbid)
monkeypatch.setattr(cls, "save_to", forbid)
monkeypatch.setattr(torch.Tensor, "detach", detach)
monkeypatch.setattr(torch.Tensor, "cpu", cpu)
actual = nodes.encode_value("VIDEO", source)
assert len(staged) == len(tensors)
assert all(id(tensor) in detached for tensor in tensors)
assert all(tensor.requires_grad for tensor in tensors)
assert actual == expected
decoded = nodes.decode_value("VIDEO", *actual)
assert type(decoded) is type(source)
assert codec.encode_video(decoded) == expected
assert len(staged) == len(tensors) # Pure codec never stages implicitly.
@pytest.mark.parametrize("field", ["images", "alpha", "audio", "complete_audio"])
def test_pure_video_rejects_non_cpu_storage_without_transfer(field, monkeypatch):
if field == "complete_audio":
source = VideoFromList([file_value()], complete_audio=audio())
target = source.complete_audio["waveform"]
else:
source = components()
target = getattr(vars(source)["_VideoFromComponents__components"], field)
if field == "audio":
target = target["waveform"]
# A property spy on real CPU tensors exercises the guard without a GPU
# allocation, subclass, device discovery or CUDA initialization.
real_device = torch.Tensor.device
monkeypatch.setattr(torch.Tensor, "device", property(
lambda tensor: types.SimpleNamespace(type="cuda") if tensor is target
else real_device.__get__(tensor)))
monkeypatch.setattr(torch.Tensor, "cpu", forbid)
with pytest.raises(ValueError, match="must be on CPU"):
codec.encode_video(source)
@pytest.mark.parametrize("field,bad", [
("_VideoFromFile__start_time", True), ("_VideoFromFile__start_time", float("nan")),
("_VideoFromFile__duration", float("inf")), ("_VideoFromFile__crop", (1, 2, True, 4)),
("_VideoFromFile__crop", [1, 2, 3, 4]), ("_VideoFromFile__crop", (1, 2, 3)),
("_VideoFromFile__file", b"not a BytesIO"),
])
def test_invalid_native_file_attributes(field, bad):
value = file_value()
setattr(value, field, bad)
with pytest.raises(ValueError):
codec.encode_video(value)
def test_missing_attributes_and_invalid_component_state():
value = file_value()
del vars(value)["_VideoFromFile__duration"]
with pytest.raises(ValueError):
codec.encode_video(value)
for field, bad in [("frame_rate", 24.0), ("frame_rate", Fraction(0)),
("alpha", torch.ones(1, 1, 1, device="cpu"))]:
value = components()
setattr(vars(value)["_VideoFromComponents__components"], field, bad)
with pytest.raises(ValueError):
codec.encode_video(value)
@pytest.mark.parametrize("metadata", [
{"value": float("nan")}, {"value": float("inf")}, {1: "non-string key"},
{"value": (1, 2)}, {"value": object()}, {"large": "x" * 65536},
])
def test_invalid_metadata_rejected_both_directions(metadata):
with pytest.raises(ValueError):
codec.encode_video(components(metadata=metadata))
descriptor, payload = codec.encode_video(components())
descriptor["metadata"] = metadata
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_metadata_depth_cycles_and_exact_size():
metadata = {}
current = metadata
for _ in range(40):
current["child"] = {}
current = current["child"]
for value in (metadata,):
with pytest.raises(ValueError):
codec.encode_video(components(metadata=value))
cyclic = {}
cyclic["self"] = cyclic
with pytest.raises(ValueError):
codec.encode_video(components(metadata=cyclic))
exact = {"x": "a" * (65536 - 8)}
assert len(json.dumps(exact, separators=(",", ":")).encode()) == 65536
descriptor, payload = codec.encode_video(components(metadata=exact))
assert codec.decode_video(descriptor, payload).get_components().metadata == exact
@pytest.mark.parametrize("field,bad", [
("encoding", "other"), ("representation", "other"), ("byte_length", True),
("byte_length", -1), ("byte_length", 12345), ("start_time", False),
("duration", float("inf")), ("crop", [0, 0, True, 4]), ("crop", [1, 2, 3]),
])
def test_malformed_file_descriptors(field, bad):
descriptor, payload = codec.encode_video(file_value())
descriptor[field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
@pytest.mark.parametrize("field,bad", [("offset", True), ("offset", -1), ("offset", 1),
("length", True), ("length", -1), ("length", 2**80)])
def test_invalid_offsets_lengths(field, bad):
descriptor, payload = codec.encode_video(file_value())
descriptor["data"][field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_truncation_trailing_gaps_overlap_and_unknown_fields():
value = VideoFromList([file_value(b"123"), file_value(b"456")])
original, payload = codec.encode_video(value)
for wire in (payload[:-1], payload + b"x"):
with pytest.raises(ValueError):
codec.decode_video(original, wire)
for offset in (2, 4):
descriptor = copy.deepcopy(original)
descriptor["segments"][1]["offset"] = offset
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["surprise"] = "bad"
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["segments"][0]["descriptor"]["byte_length"] += 1
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
descriptor = copy.deepcopy(original)
descriptor["segments"][0]["descriptor"] = descriptor
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_all_descriptors_validated_before_tensor_allocation(monkeypatch):
original, payload = codec.encode_video(components())
monkeypatch.setattr(codec, "decode_payload", forbid)
cases = []
for key, value in [("shape", [2**80, 2, 3]), ("shape", [True, 2, 3]),
("dtype", "pickle"), ("layout", "BHWC")]:
descriptor = copy.deepcopy(original)
descriptor["alpha"]["descriptor"][key] = value
cases.append(descriptor)
descriptor = copy.deepcopy(original)
descriptor["audio"]["descriptor"]["sample_rate"] = True
cases.append(descriptor)
descriptor = copy.deepcopy(original)
descriptor["alpha"]["offset"] += 1
cases.append(descriptor)
for descriptor in cases:
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
@pytest.mark.parametrize("field,bad", [("numerator", True), ("denominator", True),
("denominator", 0), ("numerator", -3),
("denominator", -1001)])
def test_invalid_fraction_fields(field, bad):
descriptor, payload = codec.encode_video(components())
descriptor["frame_rate"][field] = bad
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
def test_segment_limit_and_exact_maximum(monkeypatch):
value = VideoFromList([file_value(b"x") for _ in range(64)])
descriptor, payload = codec.encode_video(value)
result = codec.decode_video(descriptor, payload)
assert len(result.videos) == 64
value.videos.append(file_value(b"x"))
with pytest.raises(ValueError):
codec.encode_video(value)
descriptor["segments"].append(copy.deepcopy(descriptor["segments"][-1]))
monkeypatch.setattr(codec, "decode_payload", forbid)
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
value.videos = []
with pytest.raises(ValueError):
codec.encode_video(value)
value.videos = [components()]
with pytest.raises(ValueError):
codec.encode_video(value)
@pytest.mark.parametrize("descriptor,payload", [(None, b""), ([], b""), ({}, bytearray()), ({}, "")])
def test_invalid_top_level_types(descriptor, payload):
with pytest.raises(ValueError):
codec.decode_video(descriptor, payload)
+35 -1
View File
@@ -90,14 +90,48 @@ class NetworkHelpersTests(unittest.TestCase):
"https://master.example.com",
)
def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
self.assertEqual(
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
"http://192.168.68.56:8188",
)
def test_build_master_url_keeps_explicit_port_in_host(self):
cfg = {"master": {"host": "192.168.68.56:8001"}}
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
self.assertEqual(
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
"http://192.168.68.56:8001",
)
def test_build_master_url_falls_back_to_server_address(self):
cfg = {"master": {"host": ""}}
cfg = {"master": {"host": "", "port": 8001}}
prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
self.assertEqual(
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
"http://127.0.0.1:8190",
)
def test_build_master_callback_url_uses_loopback_for_local_worker(self):
cfg = {"master": {"host": "192.168.68.56"}}
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
self.assertEqual(
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
"http://127.0.0.1:8001",
)
def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
cfg = {"master": {"host": "192.168.68.56"}}
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
self.assertEqual(
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
"http://192.168.68.56:8001",
)
if __name__ == "__main__":
unittest.main()
+10 -22
View File
@@ -86,27 +86,25 @@ class ParseTilesFromFormTests(unittest.TestCase):
# --- happy paths ---
def test_single_tile_returns_one_entry(self):
def test_single_tile_returns_image_and_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertEqual(len(tiles), 1)
def test_multiple_tiles_all_returned(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(len(tiles), 3)
def test_tile_image_is_pil_image(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertIsInstance(tiles[0]["image"], PILImage.Image)
def test_tile_metadata_fields_are_parsed(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
tile = tiles[0]
self.assertIsInstance(tile["image"], PILImage.Image)
self.assertEqual(tile["tile_idx"], 0)
self.assertEqual(tile["x"], 0)
self.assertEqual(tile["y"], 0)
self.assertEqual(tile["extracted_width"], 64)
self.assertEqual(tile["extracted_height"], 64)
def test_multiple_tiles_preserve_count_order_and_coordinates(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(len(tiles), 3)
for i, tile in enumerate(tiles):
self.assertEqual(tile["tile_idx"], i)
self.assertEqual(tiles[1]["x"], 64)
self.assertEqual(tiles[2]["x"], 128)
def test_padding_is_parsed_from_form(self):
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
self.assertEqual(tiles[0]["padding"], 16)
@@ -138,16 +136,6 @@ class ParseTilesFromFormTests(unittest.TestCase):
self.assertNotIn("batch_idx", tiles[0])
self.assertNotIn("global_idx", tiles[0])
def test_tile_indices_match_metadata_order(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
for i, tile in enumerate(tiles):
self.assertEqual(tile["tile_idx"], i)
def test_x_coordinates_reflect_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(tiles[1]["x"], 64)
self.assertEqual(tiles[2]["x"], 128)
# --- error cases ---
def test_missing_tiles_metadata_raises_value_error(self):
+575 -9
View File
@@ -71,6 +71,15 @@ def _collector_only_prompt():
}
def _audio_only_prompt():
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
return {
"1": {"class_type": "LoadAudio", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
}
def _delegate_prompt():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return {
@@ -217,11 +226,74 @@ class PrunePromptForWorkerTests(unittest.TestCase):
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
def test_injects_preview_audio_for_audio_only_collector(self):
result = pt.prune_prompt_for_worker(_audio_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
prompt = _linear_prompt()
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
prompt["4"]["inputs"]["audio"] = ["6", 0]
result = pt.prune_prompt_for_worker(prompt)
self.assertEqual(
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
1,
)
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
def test_preserves_image_preview_for_distributed_upscale(self):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
def test_no_preview_image_when_no_downstream(self):
result = pt.prune_prompt_for_worker(_collector_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
self.assertEqual(len(preview_nodes), 0)
def test_worker_sink_replaces_preview_image(self):
result = pt.prune_prompt_for_worker(_linear_prompt(), use_worker_sink=True)
sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
self.assertEqual(len(sinks), 1)
self.assertEqual(sinks[0]["inputs"], {"value": ["4", 0]})
self.assertFalse(any(n.get("class_type") in ("PreviewImage", "PreviewAudio") for n in result.values()))
def test_worker_sink_uses_audio_output_for_audio_only_collector(self):
result = pt.prune_prompt_for_worker(_audio_only_prompt(), use_worker_sink=True)
sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 1]}])
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
def test_worker_sink_for_distributed_upscale(self):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 0]}])
self.assertNotIn("3", result)
def test_worker_sink_ids_do_not_collide_with_original_prompt(self):
prompt = _linear_prompt()
result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
sink_ids = [k for k, n in result.items() if n.get("class_type") == pt.WORKER_SINK_CLASS]
self.assertTrue(all(k not in prompt for k in sink_ids))
def test_no_worker_sink_when_no_downstream(self):
result = pt.prune_prompt_for_worker(_collector_only_prompt(), use_worker_sink=True)
self.assertFalse(any(n.get("class_type") == pt.WORKER_SINK_CLASS for n in result.values()))
def test_unrelated_nodes_are_pruned(self):
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
@@ -242,7 +314,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
def test_upscale_node_is_treated_as_distributed(self):
prompt = {
"1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
@@ -264,8 +336,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotIn("1", result)
self.assertNotIn("2", result)
def test_removes_dangling_upstream_refs(self):
"""Collector must not retain dangling refs to pruned upstream nodes."""
def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
"""Collector must point to exactly one valid placeholder after pruning."""
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
collector_inputs = result["3"].get("inputs", {})
@@ -276,26 +348,520 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotEqual(source_id, "2")
self.assertIn(source_id, result)
self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
def test_injects_empty_image_placeholder(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 1)
placeholder_id = empty_nodes[0][0]
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
def test_audio_only_collector_does_not_get_image_placeholder(self):
prompt = _audio_only_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(empty_nodes, [])
self.assertNotIn("images", result["2"].get("inputs", {}))
self.assertNotIn("audio", result["2"].get("inputs", {}))
def test_one_placeholder_per_collector(self):
"""Two collectors → two placeholders."""
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
"10": {"class_type": "LoadImage", "inputs": {}},
"11": {"class_type": "LoadImage", "inputs": {}},
}
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 2)
def test_preserves_primitive_string_for_downstream_required_input(self):
"""Delegate-only master keeps primitive inputs needed by SaveImage."""
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["15", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertEqual(result["15"], prompt["15"])
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
def test_preserves_load_image_for_switch_alternate_required_input(self):
"""Delegate-only master keeps LoadImage inputs needed by switches."""
prompt = {
"848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}},
"862": {"class_type": "VAEDecode", "inputs": {}},
"854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}},
"850": {
"class_type": "ComfySwitchNode",
"inputs": {
"on_false": ["848", 0],
"on_true": ["854", 0],
},
},
"851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}},
}
result = pt.prepare_delegate_master_prompt(prompt, ["854"])
self.assertIn("848", result)
self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0])
self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0])
self.assertNotIn("862", result)
def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps scalar utility chains used by SaveImage."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["16", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0])
self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0])
self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0])
def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps multi-input scalar utility chains."""
join_string_multi = type(
"JoinStringMulti",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"string_1": ("STRING",)},
"optional": {
"string_2": ("STRING",),
"string_3": ("STRING",),
},
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"JoinStringMulti": join_string_multi,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}},
"18": {
"class_type": "JoinStringMulti",
"inputs": {
"string_1": ["15", 0],
"string_2": ["16", 0],
"string_3": ["17", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["18", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertIn("18", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0])
def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self):
"""Scalar utility nodes are retained only when their full input branch is safe."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["8", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("8", result)
self.assertNotIn("17", result)
self.assertNotIn("filename_prefix", result["9"]["inputs"])
def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self):
"""Scalar outputs are retained only for scalar/config downstream inputs."""
string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)})
image_consumer = type(
"ImageConsumer",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"mask": ("IMAGE",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"ImageConsumer": image_consumer,
"StringProvider": string_provider,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"17": {"class_type": "StringProvider", "inputs": {}},
"9": {
"class_type": "ImageConsumer",
"inputs": {
"images": ["11", 0],
"mask": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("17", result)
self.assertNotIn("mask", result["9"]["inputs"])
def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps list-shaped scalar config chains."""
create_list = type(
"CreateList",
(),
{
"RETURN_TYPES": ("LIST",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"inputs.input0": ("STRING",)},
"optional": {
"inputs.input1": ("STRING",),
"inputs.input2": ("STRING",),
},
}
),
},
)
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("LIST",),
"delimiter": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"CreateList": create_list,
"SaveImage": save_image,
"StringDataListJoin": string_data_list_join,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["15", 0],
"inputs.input1": ["17", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"delimiter": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("17", result)
self.assertIn("28", result)
self.assertIn("29", result)
self.assertIn("32", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0])
self.assertNotIn("8", result)
def test_preserves_builtin_create_list_string_data_list_join_subgraph(self):
"""Delegate-only master handles ComfyUI 0.23 CreateList data-list output."""
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_IS_LIST": True,
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("STRING", {"forceInput": True}),
"sep": ("STRING", {"default": " "}),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"Basic data handling: StringDataListJoin": string_data_list_join,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["17", 0],
"inputs.input1": ["15", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "Basic data handling: StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"sep": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
for node_id in ("15", "17", "28", "29", "32"):
self.assertIn(node_id, result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0])
self.assertNotIn("8", result)
def test_result_is_independent_copy(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+167
View File
@@ -0,0 +1,167 @@
import importlib.util
import json
import sys
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def load_module():
spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py")
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
client = load_module()
def workflow_for_request_size(target_size, timeout_seconds):
"""Build an ASCII workflow whose compact UTF-8 request body has exact size."""
workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}}
def body_size(candidate):
payload = {
"input": {"workflow": candidate, "images": []},
"policy": {
"executionTimeout": timeout_seconds * 1000,
"ttl": timeout_seconds * 1000,
},
}
return len(json.dumps(
payload, allow_nan=False, ensure_ascii=False, separators=(",", ":")
).encode("utf-8"))
padding = target_size - body_size(workflow)
if padding < 0:
raise AssertionError("target is smaller than the request fixture")
workflow["1"]["inputs"]["value"] = "x" * padding
if body_size(workflow) != target_size:
raise AssertionError("request fixture did not reach the exact target size")
return workflow
class Response:
def __init__(self, status=200, payload=None):
self.status, self.payload = status, payload
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return False
async def json(self):
return self.payload
class Session:
def __init__(self, responses):
self.responses, self.calls = list(responses), []
def request(self, method, url, **kwargs):
self.calls.append((method, url, kwargs))
response = self.responses.pop(0)
if isinstance(response, Exception):
raise response
return response
class ClientTests(unittest.IsolatedAsyncioTestCase):
async def test_submit_exact_generic_contract_and_once(self):
session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})])
api = client.RunpodClient(session, "endpoint123", "secret-value")
result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800)
self.assertEqual(result, "job-1")
method, url, kwargs = session.calls[0]
self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run"))
self.assertIsInstance(kwargs["data"], bytes)
sent_payload = json.loads(kwargs["data"])
self.assertEqual(set(sent_payload["input"]), {"workflow", "images"})
self.assertEqual(sent_payload["policy"]["ttl"], 1800000)
self.assertFalse(kwargs["allow_redirects"])
self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value")
async def test_submit_timeout_is_ambiguous_and_never_retried(self):
session = Session([TimeoutError("secret-value in unsafe error")])
api = client.RunpodClient(session, "endpoint123", "secret-value")
with self.assertRaises(client.AmbiguousSubmission) as ctx:
await api.submit({}, [], 30)
self.assertNotIn("secret-value", str(ctx.exception))
self.assertEqual(len(session.calls), 1)
async def test_http_submit_failure_not_retried_or_leaked(self):
for status in [401, 429, 500, 302]:
session = Session([Response(status, {"error": "secret-value"})])
api = client.RunpodClient(session, "endpoint123", "secret-value")
with self.assertRaises(client.RunpodError) as ctx:
await api.submit({}, [], 30)
self.assertEqual(len(session.calls), 1)
self.assertNotIn("secret-value", str(ctx.exception))
async def test_health_zero_workers_is_valid(self):
session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})])
api = client.RunpodClient(session, "endpoint123", "key")
self.assertEqual((await api.health())["workers"]["idle"], 0)
async def test_invalid_provider_id_no_network(self):
session = Session([])
with self.assertRaises(ValueError):
client.RunpodClient(session, "https://attacker.invalid", "key")
api = client.RunpodClient(session, "endpoint123", "key")
with self.assertRaises(ValueError):
await api.status("../other")
self.assertFalse(session.calls)
async def test_status_and_cancel_only_target_exact_job(self):
session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})])
api = client.RunpodClient(session, "endpoint123", "key")
await api.status("job-1")
await api.cancel("job-1")
self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")])
async def test_missing_id_is_ambiguous(self):
api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key")
with self.assertRaises(client.AmbiguousSubmission):
await api.submit({}, [], 30)
async def test_large_payload_rejected_before_submit(self):
session = Session([])
api = client.RunpodClient(session, "endpoint123", "key")
with self.assertRaises(ValueError):
await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30)
self.assertFalse(session.calls)
async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self):
session = Session([Response(payload={"id": "boundary-job"})])
api = client.RunpodClient(session, "endpoint123", "key")
timeout_seconds = 37
workflow = workflow_for_request_size(10_000_000, timeout_seconds)
result = await api.submit(workflow, [], timeout_seconds)
self.assertEqual(result, "boundary-job")
sent = session.calls[0][2]["data"]
self.assertIsInstance(sent, bytes)
self.assertEqual(len(sent), 10_000_000)
self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000)
async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self):
session = Session([Response(payload={"id": "must-not-be-used"})])
api = client.RunpodClient(session, "endpoint123", "key")
workflow = workflow_for_request_size(10_000_001, 37)
workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:]
with self.assertRaisesRegex(ValueError, "10,000,001") as ctx:
await api.submit(workflow, [], 37)
self.assertNotIn("private graph value", str(ctx.exception))
self.assertFalse(session.calls)
async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self):
session = Session([Response(payload={"id": "unicode-job"})])
api = client.RunpodClient(session, "endpoint123", "key")
workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}}
await api.submit(workflow, [], 41)
sent = session.calls[0][2]["data"]
expected = json.dumps({
"input": {"workflow": workflow, "images": []},
"policy": {"executionTimeout": 41_000, "ttl": 41_000},
}, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
self.assertEqual(sent, expected)
self.assertGreater(len(sent), len(sent.decode("utf-8")))
+202
View File
@@ -0,0 +1,202 @@
import asyncio
import pytest
import torch
from test_collector_list_inputs import _load_collector_module
class _FakeManager:
def __init__(self, *, pending=False, owned=True):
self.pending = pending
self.owned = owned
self.pending_calls = []
self.owned_calls = []
self.cancel_calls = []
def is_pending(self, collector_job_id, participant_id):
self.pending_calls.append((collector_job_id, participant_id))
return self.pending
def owns(self, collector_job_id, participant_id):
self.owned_calls.append((collector_job_id, participant_id))
return self.owned
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
self.cancel_calls.append({
"worker_id": worker_id,
"execution_id": execution_id,
"collector_job_id": collector_job_id,
})
return 1
def _set_up_collector(module, job_id, queue, manager=None):
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
module.prompt_server.distributed_pending_jobs = {job_id: queue}
if manager is not None:
module.prompt_server.distributed_serverless = manager
return module.DistributedCollectorNode()
def test_serverless_pending_waits_past_normal_worker_timeout():
module = _load_collector_module()
job_id = "collector-pending"
participant = "serverless:pool:1"
queue = asyncio.Queue()
manager = _FakeManager(pending=True)
collector = _set_up_collector(module, job_id, queue, manager)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
probe_calls = []
module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
async def deliver_after_provider_work():
await asyncio.sleep(0.14)
queue.put_nowait({
"worker_id": participant,
"image_index": 0,
"tensor": torch.ones(1, 2, 2, 2),
"is_last": True,
})
async def run_collection():
producer = asyncio.create_task(deliver_after_provider_work())
result = await collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids=f'["{participant}"]',
delegate_only=True,
)
await producer
return result
images, _audio = asyncio.run(run_collection())
assert tuple(images.shape) == (1, 2, 2, 2)
assert manager.pending_calls
assert probe_calls == []
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
module = _load_collector_module()
job_id = "collector-error"
queue = asyncio.Queue()
queue.put_nowait({
"worker_id": "serverless:pool:1",
"serverless_error": "Provider execution failed safely",
"is_last": True,
})
manager = _FakeManager()
collector = _set_up_collector(module, job_id, queue, manager)
with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.cancel_calls == [{
"worker_id": None,
"execution_id": None,
"collector_job_id": job_id,
}]
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
module = _load_collector_module()
job_id = "collector-stopped"
queue = asyncio.Queue()
manager = _FakeManager(pending=False, owned=True)
collector = _set_up_collector(module, job_id, queue, manager)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.owned_calls
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_interruption_cancels_serverless_job_and_removes_collector_queue():
module = _load_collector_module()
job_id = "collector-interrupted"
queue = asyncio.Queue()
manager = _FakeManager(pending=True)
collector = _set_up_collector(module, job_id, queue, manager)
interrupted = module.comfy.model_management.InterruptProcessingException("stop")
module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
module = _load_collector_module()
job_id = "collector-shape-mismatch"
queue = asyncio.Queue()
queue.put_nowait({
"worker_id": "serverless:pool:1",
"image_index": 0,
"tensor": torch.ones(1, 3, 2, 2),
"is_last": True,
})
manager = _FakeManager()
collector = _set_up_collector(module, job_id, queue, manager)
local_images = torch.zeros(1, 2, 2, 2)
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
asyncio.run(collector.execute(
images=local_images,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
))
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
module = _load_collector_module()
job_id = "collector-regular-timeout"
queue = asyncio.Queue()
collector = _set_up_collector(module, job_id, queue)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
local_images = torch.zeros(1, 2, 2, 2)
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
images, _audio = asyncio.run(collector.execute(
images=local_images,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["worker-a"]',
))
assert torch.equal(images, local_images)
assert job_id not in module.prompt_server.distributed_pending_jobs
+436
View File
@@ -0,0 +1,436 @@
import importlib.util
import json
import os
import stat
import sys
import tempfile
import threading
import types
import unittest
from pathlib import Path
from unittest import mock
def _load_serverless_config_module(*, without_fcntl=False):
module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py"
module_name = (
"dist_serverless_config_no_fcntl_test_module"
if without_fcntl
else "dist_serverless_config_test_module"
)
package_name = module_name + "_pkg"
package = types.ModuleType(package_name)
package.__path__ = [str(module_path.parent)]
sys.modules[package_name] = package
module_name = package_name + ".serverless_config"
sys.modules.pop(module_name, None)
spec = importlib.util.spec_from_file_location(module_name, module_path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
if without_fcntl:
with mock.patch.dict(sys.modules, {"fcntl": None}):
spec.loader.exec_module(module)
else:
spec.loader.exec_module(module)
return module
serverless_config = _load_serverless_config_module()
class ValidateServerlessWorkerTests(unittest.TestCase):
def test_returns_copy_with_runpod_defaults(self):
worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}
result = serverless_config.validate_serverless_worker(worker)
self.assertEqual(
result,
{
"id": "worker-1",
"endpoint_id": "ep_123-abc",
"name": "GPU",
"type": "serverless",
"provider": "runpod",
"jobs_per_run": 1,
"timeout_seconds": 1800,
},
)
self.assertIsNot(result, worker)
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
def test_omitted_media_profile_remains_absent_and_means_legacy(self):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123"}
)
self.assertNotIn("media_profile", result)
self.assertEqual(result.get("media_profile", "legacy-image-v1"), "legacy-image-v1")
def test_explicit_media_profiles_are_retained(self):
for profile in ("legacy-image-v1", "comfy-media-v1"):
with self.subTest(profile=profile):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
)
self.assertEqual(result["media_profile"], profile)
def test_media_profile_rejects_unknown_nonstring_and_null_values(self):
for profile in ("", "auto", "comfy-media-v2", " comfy-media-v1", "COMFY-MEDIA-V1", None, True, 1, [], {}):
with self.subTest(profile=profile):
with self.assertRaisesRegex(ValueError, "media_profile"):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "media_profile": profile}
)
def test_accepts_inclusive_numeric_bounds(self):
for jobs, timeout in ((1, 30), (16, 86400)):
with self.subTest(jobs=jobs, timeout=timeout):
result = serverless_config.validate_serverless_worker(
{
"id": "worker-1",
"endpoint_id": "ep-123",
"jobs_per_run": jobs,
"timeout_seconds": timeout,
}
)
self.assertEqual(result["jobs_per_run"], jobs)
self.assertEqual(result["timeout_seconds"], timeout)
def test_legacy_concurrency_is_accepted_and_dropped(self):
# Configs saved before Images per run replaced the separate setting.
for legacy in (1, 2, 16, 17, True, "3"):
with self.subTest(legacy=legacy):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": legacy}
)
self.assertNotIn("concurrency", result)
self.assertEqual(result["jobs_per_run"], 2)
def test_endpoint_url_is_rejected_with_a_plain_message(self):
with self.assertRaisesRegex(ValueError, "not a URL"):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "https://api.runpod.ai/v2/abc123/run"}
)
def test_rejects_missing_or_empty_id_and_endpoint(self):
base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
with self.subTest(field=field, value=value):
worker = dict(base)
worker[field] = value
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
worker = dict(base)
del worker["id"]
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
def test_endpoint_id_must_be_ascii_opaque_identifier(self):
for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"):
with self.subTest(endpoint_id=endpoint_id):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": endpoint_id}
)
def test_endpoint_id_matches_runpod_client_length_limit(self):
valid_id = "e" * 200
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": valid_id}
)
self.assertEqual(result["endpoint_id"], valid_id)
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "e" * 201}
)
def test_enabled_must_be_boolean_when_supplied(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
for enabled in (True, False):
with self.subTest(enabled=enabled):
self.assertIs(
serverless_config.validate_serverless_worker(
{**base, "enabled": enabled}
)["enabled"],
enabled,
)
for enabled in ("false", 0, 1, None):
with self.subTest(enabled=enabled):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{**base, "enabled": enabled}
)
def test_rejects_unknown_provider_and_non_serverless_type(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
for extra in ({"provider": "other"}, {"type": "local"}):
with self.subTest(extra=extra):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker({**base, **extra})
def test_numeric_fields_reject_bool_float_and_out_of_range(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
invalid_values = {
"jobs_per_run": (True, 1.0, 0, 17),
"timeout_seconds": (True, 30.0, 29, 86401),
}
for field, values in invalid_values.items():
for value in values:
with self.subTest(field=field, value=value):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker({**base, field: value})
def test_requires_worker_object(self):
for worker in (None, [], "worker"):
with self.subTest(worker=worker):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
class SecretStoreTests(unittest.TestCase):
def setUp(self):
self.tempdir = tempfile.TemporaryDirectory()
self.root = Path(self.tempdir.name)
self.path = self.root / "private" / "runpod-secrets.json"
def tearDown(self):
self.tempdir.cleanup()
def test_set_get_has_delete_and_persists_as_json_mapping(self):
store = serverless_config.SecretStore(self.path)
store.set("worker-1", "secret-value")
self.assertEqual(store.get("worker-1"), "secret-value")
self.assertTrue(store.has("worker-1"))
self.assertIsNone(store.get("missing"))
self.assertFalse(store.has("missing"))
self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"})
self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value")
store.delete("worker-1")
self.assertIsNone(store.get("worker-1"))
self.assertFalse(store.has("worker-1"))
def test_rejects_api_keys_with_newlines_or_over_4096_characters(self):
store = serverless_config.SecretStore(self.path)
for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097):
with self.subTest(key_length=len(key)):
with self.assertRaises(serverless_config.SecretStoreError):
store.set("worker-1", key)
store.set("worker-1", "k" * 4096)
self.assertEqual(store.get("worker-1"), "k" * 4096)
def test_new_secret_file_and_directory_have_restrictive_modes(self):
store = serverless_config.SecretStore(self.path)
store.set("worker-1", "secret-value")
self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600)
self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700)
def test_missing_store_reads_as_empty_without_creating_files(self):
store = serverless_config.SecretStore(self.path)
self.assertIsNone(store.get("worker-1"))
self.assertFalse(store.has("worker-1"))
self.assertFalse(self.path.parent.exists())
def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self):
self.path.parent.mkdir()
corrupt = b'{"partial": "NEVER-LEAK-THIS-secret'
self.path.write_bytes(corrupt)
store = serverless_config.SecretStore(self.path)
for operation in (
lambda: store.get("worker-1"),
lambda: store.set("worker-1", "replacement-secret"),
lambda: store.delete("worker-1"),
):
with self.subTest(operation=operation):
with self.assertRaises(Exception) as raised:
operation()
self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception))
self.assertNotIn("replacement-secret", str(raised.exception))
self.assertEqual(self.path.read_bytes(), corrupt)
def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self):
for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'):
with self.subTest(corrupt=corrupt):
self.path.parent.mkdir(exist_ok=True)
self.path.write_bytes(corrupt)
with self.assertRaises(Exception) as raised:
serverless_config.SecretStore(self.path).get("worker")
self.assertNotIn("opaque-stored-value", str(raised.exception))
def test_rejects_symlink_file_and_preserves_link_target(self):
target = self.root / "target.json"
target.write_text('{"worker-1": "target-secret"}', encoding="utf-8")
self.path.parent.mkdir()
self.path.symlink_to(target)
store = serverless_config.SecretStore(self.path)
with self.assertRaises(Exception):
store.get("worker-1")
with self.assertRaises(Exception):
store.set("worker-1", "replacement-secret")
self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}')
def test_rejects_symlinked_parent_directory(self):
real_dir = self.root / "real"
real_dir.mkdir()
link_dir = self.root / "linked"
link_dir.symlink_to(real_dir, target_is_directory=True)
store = serverless_config.SecretStore(link_dir / "secrets.json")
with self.assertRaises(Exception):
store.set("worker-1", "secret-value")
self.assertFalse((real_dir / "secrets.json").exists())
def test_rejects_parent_traversal_in_store_path(self):
store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json")
with self.assertRaises(Exception):
store.set("worker-1", "secret-value")
self.assertFalse((self.root / "secrets.json").exists())
def test_threaded_updates_from_independent_instances_are_not_lost(self):
workers = 12
barrier = threading.Barrier(workers)
errors = []
def write_one(index):
try:
barrier.wait(timeout=5)
serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}")
except Exception as exc: # collected to fail in the test thread
errors.append(exc)
threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=10)
self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish")
self.assertEqual(errors, [])
saved = json.loads(self.path.read_text(encoding="utf-8"))
self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)})
class PublicConfigTests(unittest.TestCase):
def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self):
with tempfile.TemporaryDirectory() as tempdir:
store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json")
store.set("serverless-1", "stored-api-key")
config = {
"workers": [
{
"id": "serverless-1",
"type": "serverless",
"provider": "runpod",
"endpoint_id": "ep-123",
"media_profile": "comfy-media-v1",
"api_key": "inline-secret",
"clear_api_key": "another-secret",
"credential": "credential-secret",
"nested": {"runpod_credential": "nested-secret", "safe": "kept"},
"name": "remote GPU",
},
{"id": "serverless-2", "type": "serverless", "name": "not configured"},
{"id": "local-1", "type": "local", "api_key": "local-secret"},
],
"settings": {"mode": "preserved", "credentials_note": "remove-me"},
}
original = json.loads(json.dumps(config))
result = serverless_config.public_config(config, store)
first, second, local = result["workers"]
self.assertEqual(first["api_key_configured"], True)
self.assertEqual(second["api_key_configured"], False)
self.assertNotIn("api_key_configured", local)
self.assertNotIn("api_key", first)
self.assertNotIn("clear_api_key", first)
self.assertNotIn("credential", first)
self.assertNotIn("runpod_credential", first["nested"])
self.assertEqual(first["nested"]["safe"], "kept")
self.assertNotIn("api_key", local)
self.assertNotIn("credentials_note", result["settings"])
self.assertEqual(first["name"], "remote GPU")
self.assertEqual(first["media_profile"], "comfy-media-v1")
self.assertNotIn("media_profile", second)
self.assertEqual(config, original)
self.assertIsNot(result, config)
self.assertIsNot(result["workers"], config["workers"])
def test_redacts_normalized_secret_field_aliases(self):
config = {
"workers": [{
"id": "serverless-1",
"type": "serverless",
"apiKey": "camel-api-key",
"APIKEY": "upper-api-key",
"clearApiKey": "camel-clear-key",
"CLEAR_API_KEY": "upper-clear-key",
"runpod_CREDENTIAL": "provider-credential",
"nested": {"credential_note": "nested-credential"},
"safe": "preserved",
}]
}
result = serverless_config.public_config(
config, serverless_config.SecretStore(self.root_path())
)
worker = result["workers"][0]
self.assertEqual(worker, {
"id": "serverless-1",
"type": "serverless",
"nested": {},
"safe": "preserved",
"api_key_configured": False,
})
def test_handles_non_dict_config_without_mutating_it(self):
config = ["safe", {"api_key": "secret"}]
result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path()))
self.assertEqual(result, ["safe", {}])
self.assertEqual(config, ["safe", {"api_key": "secret"}])
def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self):
module = _load_serverless_config_module(without_fcntl=True)
store = module.SecretStore(self.root_path())
config = {
"workers": [
{"id": "local-1", "type": "local", "api_key": "inline-secret"}
]
}
result = module.public_config(config, store)
self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]})
self.assertEqual(config["workers"][0]["api_key"], "inline-secret")
def test_secret_operations_fail_clearly_without_fcntl(self):
module = _load_serverless_config_module(without_fcntl=True)
store = module.SecretStore(self.root_path())
unsupported_message = "Serverless credential storage is currently supported on POSIX masters only"
for operation in (
lambda: store.get("worker-1"),
lambda: store.set("worker-1", "api-key"),
lambda: store.delete("worker-1"),
lambda: store.has("worker-1"),
):
with self.subTest(operation=operation):
with self.assertRaisesRegex(module.SecretStoreError, unsupported_message):
operation()
@staticmethod
def root_path():
return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json"
if __name__ == "__main__":
unittest.main()

Some files were not shown because too many files have changed in this diff Show More