fix: harden per-GPU preview outputs after review

- Typed serverless pools validate the unpruned workflow, so a preview on
  any new slot failed the whole run; both serverless compilers now accept
  every Collector output slot (COLLECTOR_OUTPUT_COUNT).
- Serverless pools are config workers: their serverless:<pool>:<job>
  participants now share the pool's worker_N slot instead of overflowing.
- A worker returning only empty frames no longer raises in the preview
  split (the combined batch already skipped them).
- Worker previews are views into the combined batch, not a second copy,
  and a frame-count guard skips them rather than misattribute frames.
- Frontend: one shared config fetch per load; unnamed slots drop their
  saved label instead of storing a redundant one.
- Tooltip notes that the per-GPU outputs are unavailable downstream of
  Ultimate SD Upscale Distributed.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
speedyrulz
2026-10-31 16:42:40 -04:00
co-authored by Claude Fable 5.1
parent 18ff755001
commit 339719ece7
8 changed files with 217 additions and 50 deletions
+65
View File
@@ -283,3 +283,68 @@ def test_undistributed_run_previews_everything_as_master(media_modules):
pass_through = module.DistributedCollectorNode()
pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True])
assert pass_through.participant_images == {}
def test_empty_worker_batch_keeps_master_preview_and_combined_result(media_modules):
module, _ = media_modules
module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait({"worker_id": "worker-a", "image_index": 0,
"tensor": torch.empty(0, 8, 12, 3), "is_last": True})
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
master = torch.zeros(1, 8, 12, 3)
collector = module.DistributedCollectorNode()
combined, _ = asyncio.run(collector.execute(
master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]'))
assert combined.shape[0] == 1
assert list(collector.participant_images) == ["master"]
def test_serverless_pool_jobs_share_their_pools_preview_slot(media_modules):
module, _ = media_modules
module.load_config = lambda: {"workers": [{"id": "gpu-a"}, {"id": "pool", "type": "serverless"}]}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
for job, value in ((2, 0.75), (1, 0.5)):
queue.put_nowait({"worker_id": f"serverless:pool:{job}", "image_index": 0,
"tensor": torch.full((1, 8, 12, 3), value), "is_last": True})
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
collector = module.DistributedCollectorNode()
asyncio.run(collector.execute(
torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test",
enabled_worker_ids='["serverless:pool:1", "serverless:pool:2"]'))
previews = collector.participant_images
assert set(previews) == {"master", "worker_2"}
# Pool jobs appear in their enabled order, not delivery order.
assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.5, 0.75]
def test_serverless_validators_accept_every_collector_output(media_modules):
module, _ = media_modules
legacy = _load_file("serverless_prompt_slot_check", ROOT / "api/orchestration/serverless_prompt.py")
assert legacy.COLLECTOR_OUTPUT_COUNT == 2 + len(module.PER_GPU_OUTPUT_NAMES)
assert legacy._REMOVED_NODE_OUTPUTS["DistributedCollector"] == set(range(legacy.COLLECTOR_OUTPUT_COUNT))
def test_worker_previews_are_views_of_the_combined_batch(media_modules):
module, _ = media_modules
module.load_config = lambda: {"workers": [{"id": "worker-a"}]}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
for index in range(2):
queue.put_nowait({"worker_id": "worker-a", "image_index": index,
"tensor": torch.full((1, 8, 12, 3), 0.5), "is_last": index == 1})
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
collector = module.DistributedCollectorNode()
combined, _ = asyncio.run(collector.execute(
torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]'))
preview = collector.participant_images["worker_1"]
assert preview.shape == (2, 8, 12, 3) and preview.is_contiguous()
assert preview.data_ptr() == combined[1:].data_ptr()
+18
View File
@@ -424,3 +424,21 @@ assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1"
'''
result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True)
assert result.returncode == 0, result.stderr
def test_links_from_collector_preview_outputs_are_pruned_not_rejected():
# Per-GPU preview sockets follow images/audio; a pool never returns them,
# but a workflow previewing them must still compile for typed pools.
prompt = _prompt()
prompt["preview:1"] = _node("PreviewImage", images=["collector:1", 2])
prompt["preview:2"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT - 1])
graph, bindings = _run(prompt)
assert [binding["socket"] for binding in bindings] == ["images"]
assert not any(node["class_type"] == "PreviewImage" for node in graph.values())
def test_rejects_collector_output_slots_beyond_the_schema():
prompt = _prompt()
prompt["preview:1"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT])
with pytest.raises(ValueError, match="invalid output slot"):
_run(prompt)