Files
robertvoy-ComfyUI-Distributed/tests/test_prompt_transform.py

494 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import importlib.util
import json
import sys
import types
import unittest
from pathlib import Path
def _load_prompt_transform_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
package_name = "dist_pt_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
api_pkg = types.ModuleType(f"{package_name}.api")
api_pkg.__path__ = []
sys.modules[f"{package_name}.api"] = api_pkg
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
orch_pkg.__path__ = []
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{package_name}.utils"] = utils_pkg
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
spec = importlib.util.spec_from_file_location(
f"{package_name}.api.orchestration.prompt_transform",
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
pt = _load_prompt_transform_module()
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _linear_prompt():
"""1 → 2 → 3 → 4(DistributedCollector) → 5(SaveImage)"""
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1]}},
"3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
"5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
}
def _collector_only_prompt():
"""1(Checkpoint) → 2(DistributedCollector) [no downstream from 2]"""
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
}
def _delegate_prompt():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}},
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
"4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
}
def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False):
if enabled_worker_ids is None:
enabled_worker_ids = ["worker-a", "worker-b"]
idx = pt.PromptIndex(prompt)
job_id_map = pt.generate_job_id_map(idx, "run")
return pt.apply_participant_overrides(
prompt,
participant_id=participant_id,
enabled_worker_ids=enabled_worker_ids,
job_id_map=job_id_map,
master_url="http://master.example.com",
delegate_master=delegate_master,
prompt_index=idx,
)
# ---------------------------------------------------------------------------
# PromptIndex
# ---------------------------------------------------------------------------
class PromptIndexTests(unittest.TestCase):
def test_nodes_by_class_groups_correctly(self):
prompt = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
"3": {"class_type": "DistributedCollector", "inputs": {}},
}
idx = pt.PromptIndex(prompt)
self.assertCountEqual(idx.nodes_for_class("DistributedCollector"), ["2", "3"])
self.assertEqual(idx.nodes_for_class("CheckpointLoaderSimple"), ["1"])
def test_nodes_for_class_unknown_returns_empty(self):
idx = pt.PromptIndex({"1": {"class_type": "KSampler", "inputs": {}}})
self.assertEqual(idx.nodes_for_class("Nonexistent"), [])
def test_nodes_without_class_type_are_indexed_under_none(self):
prompt = {"1": {"inputs": {}}}
idx = pt.PromptIndex(prompt)
# Should not raise; nodes_for_class with None key or missing class_type
self.assertEqual(idx.nodes_for_class("KSampler"), [])
def test_copy_prompt_is_a_deep_copy(self):
prompt = {"1": {"class_type": "KSampler", "inputs": {"seed": 42}}}
idx = pt.PromptIndex(prompt)
copy = idx.copy_prompt()
copy["1"]["inputs"]["seed"] = 999
self.assertEqual(prompt["1"]["inputs"]["seed"], 42)
def test_has_upstream_direct_connection(self):
"""Node 4 reads directly from node 3 (KSampler)."""
idx = pt.PromptIndex(_linear_prompt())
self.assertTrue(idx.has_upstream("4", "KSampler"))
def test_has_upstream_transitive_connection(self):
"""Node 4 → 3 → 2 → 1 (CheckpointLoaderSimple)."""
idx = pt.PromptIndex(_linear_prompt())
self.assertTrue(idx.has_upstream("4", "CheckpointLoaderSimple"))
def test_has_upstream_returns_false_when_no_path(self):
idx = pt.PromptIndex(_linear_prompt())
# CheckpointLoaderSimple has no upstream nodes
self.assertFalse(idx.has_upstream("1", "DistributedCollector"))
def test_has_upstream_result_is_cached(self):
idx = pt.PromptIndex(_linear_prompt())
r1 = idx.has_upstream("4", "KSampler")
r2 = idx.has_upstream("4", "KSampler")
self.assertEqual(r1, r2)
self.assertIn(("4", "KSampler"), idx._upstream_cache)
def test_has_upstream_does_not_infinite_loop_on_cycle(self):
"""Cyclic references in inputs should not cause infinite recursion."""
prompt = {
"1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
"2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
}
idx = pt.PromptIndex(prompt)
# Should terminate without error
result = idx.has_upstream("1", "NonExistent")
self.assertFalse(result)
# ---------------------------------------------------------------------------
# find_nodes_by_class
# ---------------------------------------------------------------------------
class FindNodesByClassTests(unittest.TestCase):
def test_finds_matching_nodes(self):
prompt = {
"1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
}
result = pt.find_nodes_by_class(prompt, "KSampler")
self.assertEqual(result, ["1"])
def test_returns_empty_when_no_match(self):
prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
self.assertEqual(pt.find_nodes_by_class(prompt, "DistributedCollector"), [])
def test_skips_non_dict_nodes(self):
prompt = {"1": "not a dict", "2": {"class_type": "KSampler", "inputs": {}}}
result = pt.find_nodes_by_class(prompt, "KSampler")
self.assertEqual(result, ["2"])
# ---------------------------------------------------------------------------
# prune_prompt_for_worker
# ---------------------------------------------------------------------------
class PrunePromptForWorkerTests(unittest.TestCase):
def test_no_distributed_nodes_returns_prompt_unchanged(self):
prompt = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
self.assertCountEqual(result.keys(), ["1", "2"])
def test_keeps_collector_and_upstream(self):
prompt = _linear_prompt()
result = pt.prune_prompt_for_worker(prompt)
for node_id in ("1", "2", "3", "4"):
self.assertIn(node_id, result)
def test_removes_downstream_of_collector(self):
prompt = _linear_prompt()
result = pt.prune_prompt_for_worker(prompt)
self.assertNotIn("5", result)
def test_injects_preview_image_when_downstream_exists(self):
prompt = _linear_prompt()
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"], ["4", 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_unrelated_nodes_are_pruned(self):
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "UnrelatedNode", "inputs": {}}, # no connection to 1
}
result = pt.prune_prompt_for_worker(prompt)
self.assertIn("1", result)
self.assertNotIn("2", result)
def test_result_is_a_copy_not_same_object(self):
prompt = _linear_prompt()
result = pt.prune_prompt_for_worker(prompt)
# Mutating the result should not affect the original
original_keys = set(prompt.keys())
result["NEW"] = {"class_type": "Test", "inputs": {}}
self.assertEqual(set(prompt.keys()), original_keys)
def test_upscale_node_is_treated_as_distributed(self):
prompt = {
"1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
self.assertIn("1", result)
self.assertIn("2", result)
self.assertNotIn("3", result)
# ---------------------------------------------------------------------------
# prepare_delegate_master_prompt
# ---------------------------------------------------------------------------
class PrepareDelegateMasterPromptTests(unittest.TestCase):
def test_keeps_collector_and_downstream(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
self.assertIn("3", result)
self.assertIn("4", result)
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."""
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
collector_inputs = result["3"].get("inputs", {})
# Original "images" pointed at node 2, which is pruned.
# It should now point at a newly injected placeholder node.
self.assertIn("images", collector_inputs)
source_id = str(collector_inputs["images"][0])
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_one_placeholder_per_collector(self):
"""Two collectors → two placeholders."""
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
}
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_result_is_independent_copy(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
result["3"]["inputs"]["NEW"] = "injected"
# Original should be untouched
self.assertNotIn("NEW", prompt["3"].get("inputs", {}))
# ---------------------------------------------------------------------------
# generate_job_id_map
# ---------------------------------------------------------------------------
class GenerateJobIdMapTests(unittest.TestCase):
def test_maps_collector_nodes(self):
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "KSampler", "inputs": {}},
}
idx = pt.PromptIndex(prompt)
job_map = pt.generate_job_id_map(idx, "prefix")
self.assertEqual(job_map["1"], "prefix_1")
self.assertNotIn("2", job_map)
def test_maps_upscale_nodes(self):
prompt = {
"5": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
}
idx = pt.PromptIndex(prompt)
job_map = pt.generate_job_id_map(idx, "run")
self.assertEqual(job_map["5"], "run_5")
def test_empty_prompt_returns_empty_map(self):
idx = pt.PromptIndex({})
self.assertEqual(pt.generate_job_id_map(idx, "prefix"), {})
def test_stable_ids_across_calls(self):
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
idx = pt.PromptIndex(prompt)
m1 = pt.generate_job_id_map(idx, "run")
m2 = pt.generate_job_id_map(idx, "run")
self.assertEqual(m1, m2)
# ---------------------------------------------------------------------------
# apply_participant_overrides – DistributedCollector
# ---------------------------------------------------------------------------
class ApplyOverridesCollectorTests(unittest.TestCase):
def _collector_prompt(self):
return {"1": {"class_type": "DistributedCollector", "inputs": {}}}
def test_worker_sets_is_worker_true(self):
result = _apply(self._collector_prompt(), "worker-a")
self.assertTrue(result["1"]["inputs"]["is_worker"])
def test_worker_sets_master_url(self):
result = _apply(self._collector_prompt(), "worker-a")
self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
def test_worker_sets_worker_id(self):
result = _apply(self._collector_prompt(), "worker-a")
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
def test_worker_sets_delegate_only_false(self):
result = _apply(self._collector_prompt(), "worker-a")
self.assertFalse(result["1"]["inputs"]["delegate_only"])
def test_master_sets_is_worker_false(self):
result = _apply(self._collector_prompt(), "master")
self.assertFalse(result["1"]["inputs"]["is_worker"])
def test_master_clears_stale_master_url(self):
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"master_url": "stale"}}}
result = _apply(prompt, "master")
self.assertNotIn("master_url", result["1"]["inputs"])
def test_master_clears_stale_worker_id(self):
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {"worker_id": "stale"}}}
result = _apply(prompt, "master")
self.assertNotIn("worker_id", result["1"]["inputs"])
def test_master_with_delegate_master_sets_delegate_only_true(self):
result = _apply(self._collector_prompt(), "master", delegate_master=True)
self.assertTrue(result["1"]["inputs"]["delegate_only"])
def test_master_without_delegate_master_sets_delegate_only_false(self):
result = _apply(self._collector_prompt(), "master", delegate_master=False)
self.assertFalse(result["1"]["inputs"]["delegate_only"])
def test_enabled_worker_ids_serialized_as_json(self):
enabled = ["worker-a", "worker-b"]
result = _apply(self._collector_prompt(), "master", enabled_worker_ids=enabled)
self.assertEqual(result["1"]["inputs"]["enabled_worker_ids"], json.dumps(enabled))
def test_multi_job_id_is_set_from_job_map(self):
prompt = {"1": {"class_type": "DistributedCollector", "inputs": {}}}
idx = pt.PromptIndex(prompt)
job_id_map = {"1": "run_abc_1"}
result = pt.apply_participant_overrides(
prompt,
participant_id="worker-a",
enabled_worker_ids=["worker-a"],
job_id_map=job_id_map,
master_url="http://master",
delegate_master=False,
prompt_index=idx,
)
self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_abc_1")
# ---------------------------------------------------------------------------
# apply_participant_overrides – DistributedSeed
# ---------------------------------------------------------------------------
class ApplyOverridesSeedTests(unittest.TestCase):
def _seed_prompt(self):
return {"1": {"class_type": "DistributedSeed", "inputs": {}}}
def test_worker_sets_is_worker_true(self):
result = _apply(self._seed_prompt(), "worker-a")
self.assertTrue(result["1"]["inputs"]["is_worker"])
def test_worker_id_reflects_index_in_enabled_list(self):
result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
def test_master_sets_is_worker_false(self):
result = _apply(self._seed_prompt(), "master")
self.assertFalse(result["1"]["inputs"]["is_worker"])
def test_master_sets_empty_worker_id(self):
result = _apply(self._seed_prompt(), "master")
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
# ---------------------------------------------------------------------------
# apply_participant_overrides – UltimateSDUpscaleDistributed
# ---------------------------------------------------------------------------
class ApplyOverridesUpscaleTests(unittest.TestCase):
def _upscale_prompt(self):
return {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}}}
def test_worker_sets_is_worker_true(self):
result = _apply(self._upscale_prompt(), "worker-a")
self.assertTrue(result["1"]["inputs"]["is_worker"])
def test_worker_sets_master_url_and_worker_id(self):
result = _apply(self._upscale_prompt(), "worker-a")
self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
def test_master_clears_master_url_and_worker_id(self):
prompt = {"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"master_url": "x", "worker_id": "y"}}}
result = _apply(prompt, "master")
self.assertNotIn("master_url", result["1"]["inputs"])
self.assertNotIn("worker_id", result["1"]["inputs"])
def test_collector_downstream_of_upscale_gets_pass_through(self):
"""A DistributedCollector that is downstream of UltimateSDUpscaleDistributed → pass_through=True."""
prompt = {
"1": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
}
result = _apply(prompt, "worker-a", enabled_worker_ids=["worker-a"])
self.assertTrue(result["2"]["inputs"].get("pass_through"))
# ---------------------------------------------------------------------------
# apply_participant_overrides – DistributedValue
# ---------------------------------------------------------------------------
class ApplyOverridesValueTests(unittest.TestCase):
def _value_prompt(self):
return {"1": {"class_type": "DistributedValue", "inputs": {}}}
def test_worker_sets_is_worker_true(self):
result = _apply(self._value_prompt(), "worker-a")
self.assertTrue(result["1"]["inputs"]["is_worker"])
def test_worker_id_reflects_index_in_enabled_list(self):
result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
def test_master_sets_is_worker_false(self):
result = _apply(self._value_prompt(), "master")
self.assertFalse(result["1"]["inputs"]["is_worker"])
def test_master_sets_empty_worker_id(self):
result = _apply(self._value_prompt(), "master")
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
if __name__ == "__main__":
unittest.main()