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:
co-authored by
Claude Fable 5.1
parent
18ff755001
commit
339719ece7
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user