import importlib.util import json import sys import types from pathlib import Path import pytest MODULE_PATH = ( Path(__file__).resolve().parents[1] / "api" / "orchestration" / "serverless_prompt.py" ) _spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH) assert _spec is not None and _spec.loader is not None _serverless_prompt = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_serverless_prompt) compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt def _prepared_prompt(): return { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}}, "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}}, "3": { "class_type": "DistributedSeed", "inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"}, }, "4": { "class_type": "DistributedCollector", "inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True}, }, "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}}, "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}}, "7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}}, "8": {"class_type": "UnrelatedNode", "inputs": {}}, } def _run(prompt): return compile_serverless_prompt(prompt) def _load_prompt_transform(): module_path = MODULE_PATH.with_name("prompt_transform.py") package_name = "serverless_prompt_transform_testpkg" root_pkg = types.ModuleType(package_name) root_pkg.__path__ = [] sys.modules[package_name] = root_pkg for suffix in ("api", "api.orchestration", "utils"): pkg = types.ModuleType(f"{package_name}.{suffix}") pkg.__path__ = [] sys.modules[f"{package_name}.{suffix}"] = pkg 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) 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 def test_compiles_prompt_prepared_by_existing_worker_helpers(): prompt_transform = _load_prompt_transform() original = { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}}, "3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}}, } pruned = prompt_transform.prune_prompt_for_worker(original) prompt_index = prompt_transform.PromptIndex(pruned) prepared = prompt_transform.apply_participant_overrides( pruned, participant_id="worker-a", enabled_worker_ids=["worker-a"], job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"), master_url="http://master.invalid", delegate_master=False, prompt_index=prompt_index, ) graph, manifest = _run(prepared) assert graph["2"]["inputs"]["seed"] == 101 assert set(graph) == {"1", "2", "7"} assert graph["7"]["class_type"] == "SaveImage" assert manifest[0]["job_id"] == "run_4" assert manifest[0]["collector_id"] == "4" def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream(): original = _prepared_prompt() graph, manifest = _run(original) assert manifest == [ { "collector_id": "4", "job_id": "job/42", "node_id": "9", "kind": "image", "prefix": graph["9"]["inputs"]["filename_prefix"], } ] prefix = manifest[0]["prefix"] assert graph["9"] == { "class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": prefix}, "_meta": {"title": "Serverless output: collector 4 (image)"}, } assert "/" not in prefix and "\\" not in prefix assert graph["2"]["inputs"]["seed"] == 41 assert set(graph) == {"1", "2", "9"} assert all("Distributed" not in node["class_type"] for node in graph.values()) assert original["3"]["inputs"]["seed"] == 40 assert original["4"]["inputs"]["images"] == ["2", 0] def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts(): prompt = _prepared_prompt() prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe") first_graph, first_manifest = _run(prompt) second_graph, second_manifest = _run(prompt) assert first_manifest == second_manifest assert first_graph == second_graph prefix = first_manifest[0]["prefix"] assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"] assert not any(char in prefix for char in ("/", "\\", ":")) assert first_manifest[0]["job_id"] == "../run/a" def test_compiles_multiple_disconnected_collectors_in_stable_order(): prompt = { "10": {"class_type": "LoadImage", "inputs": {}}, "20": {"class_type": "LoadImage", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}}, "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}}, "30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}}, } graph, manifest = _run(prompt) assert [item["collector_id"] for item in manifest] == ["1", "2"] assert [item["node_id"] for item in manifest] == ["31", "32"] assert graph["31"]["inputs"]["images"] == ["10", 0] assert graph["32"]["inputs"]["images"] == ["20", 0] assert set(graph) == {"10", "20", "31", "32"} def test_rejects_audio_input_instead_of_emitting_audio_saver(): prompt = { "1": {"class_type": "LoadAudio", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}}, } with pytest.raises(ValueError, match="audio.*unsupported"): _run(prompt) def test_rejects_collector_without_image_input(): with pytest.raises(ValueError, match="collector.*image"): _run({"1": {"class_type": "DistributedCollector", "inputs": {}}}) @pytest.mark.parametrize( "class_type", ["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"], ) def test_rejects_unsupported_distributed_and_batch_nodes(class_type): prompt = { "1": {"class_type": class_type, "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, "3": {"class_type": "LoadImage", "inputs": {}}, } with pytest.raises(ValueError, match="unsupported"): _run(prompt) def test_rejects_collector_to_collector_dependency_through_processing_node(): prompt = { "1": {"class_type": "LoadImage", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}}, "3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, } with pytest.raises(ValueError, match="collector-to-collector"): _run(prompt) @pytest.mark.parametrize( "link", [["missing", 0], ["1", -1], ["1", "0"], ["1", True]], ) def test_rejects_malformed_api_links(link): prompt = { "1": {"class_type": "LoadImage", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": link}}, } with pytest.raises(ValueError, match="link"): _run(prompt) def test_rejects_unexpected_output_slot_from_removed_distributed_node(): prompt = { "1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}}, "2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}}, "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, } with pytest.raises(ValueError, match="output slot"): _run(prompt) def test_distributed_seed_matches_worker_offset_and_master_passthrough(): prompt = { "1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}}, "2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}}, "3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, } graph, _manifest = _run(prompt) assert graph["3"]["inputs"] == {"seed": 12, "steps": 7} assert "1" not in graph and "2" not in graph def test_distributed_value_uses_one_based_worker_values_and_type_coercion(): prompt = { "1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}}, "2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}}, "3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, } graph, _manifest = _run(prompt) assert graph["3"]["inputs"] == {"steps": 8, "seed": 3} assert "1" not in graph and "2" not in graph def test_resolves_stock_primitive_links_for_distributed_seed_and_value(): prompt = { "1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}}, "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}}, "3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}}, "4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}}, "5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}}, "6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}}, "7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}}, } graph, _manifest = _run(prompt) assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75} assert not {"1", "2", "3", "4", "5"}.intersection(graph) def test_rejects_dynamic_linked_seed_or_value_without_executing_it(): prompt = { "1": {"class_type": "ArbitraryNode", "inputs": {}}, "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}}, "3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, } with pytest.raises(ValueError, match="literal or supported Primitive"): _run(prompt) def test_primitive_resolution_detects_cycles(): prompt = { "1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}}, "2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}}, "3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}}, "4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}}, "5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}}, } with pytest.raises(ValueError, match="cycle"): _run(prompt) def test_model_name_trivial_text_passthrough_is_inlined(): prompt = { "1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}}, "2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}}, "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, } graph, _manifest = _run(prompt) assert graph["2"]["inputs"]["model_name"] == "model.safetensors" assert "1" not in graph def test_rejects_dynamic_linked_model_name_text(): prompt = { "1": {"class_type": "ArbitraryNode", "inputs": {}}, "2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}}, "3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}}, "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, } with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"): _run(prompt) def test_resolves_boolean_and_float_primitive_links(): prompt = { "1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}}, "2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}, "3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}}, "4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}}, "5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}}, "6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}}, } graph, _manifest = _run(prompt) assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5} assert not {"1", "2", "3", "4"}.intersection(graph) def test_compiled_graph_matches_golden_prompt_and_manifest(): prompt = _prepared_prompt() prompt["3"]["inputs"]["seed"] = ["10", 0] prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}} prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}} prompt["12"] = { "class_type": "DistributedValue", "inputs": { "default_value": ["11", 0], "worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}), "is_worker": True, "worker_id": "worker_0", }, } prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}} prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}} prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0]) graph, manifest = _run(prompt) expected_prefix = ( "serverless_job_42_worker_0_4_image_" "1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244" ) assert list(graph) == ["1", "2", "15"] assert graph == { "1": { "class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}, }, "2": { "class_type": "KSampler", "inputs": { "model": ["1", 0], "seed": 41, "cfg": 2.5, "model_name": "model.safetensors", }, }, "15": { "class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": expected_prefix}, "_meta": {"title": "Serverless output: collector 4 (image)"}, }, } assert manifest == [ { "collector_id": "4", "job_id": "job/42", "node_id": "15", "kind": "image", "prefix": expected_prefix, } ]