nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
ttulttul-ComfyUI-Modal/tests/test_mapped_execution.py
T

2535 lines
92 KiB
Python
Raw Permalink Blame History

"""Tests split from the Modal executor integration suite."""
from __future__ import annotations
from modal_executor_test_support import * # noqa: F401,F403
def test_mapped_boundary_payload_requests_scheduler_list_output(
api_intercept_module: Any,
component_planning_module: Any,
) -> None:
"""Explicit mapped outputs should tell the local proxy to expose ComfyUI list semantics."""
boundary_output = api_intercept_module.BoundaryOutputSpec(
proxy_output_name="sampler_samples",
source=api_intercept_module.LinkedOutputRef(node_id="12", output_index=0),
io_type="LATENT",
is_list=False,
)
payload = component_planning_module._boundary_output_payload(
boundary_output,
mapped_output=True,
)
assert payload["mapped_output"] is True
assert payload["scheduler_is_list"] is True
assert component_planning_module._proxy_boundary_output_is_list(payload) is True
def test_modal_cloud_mirrors_phase_logs_to_stdout_in_modal_runtime(
modal_cloud_module: Any,
monkeypatch: Any,
capsys: Any,
) -> None:
"""Timed cloud phases should write directly to stdout inside Modal containers."""
monkeypatch.setenv("MODAL_IS_REMOTE", "1")
with modal_cloud_module._timed_phase("phase_under_test", component="component-1"):
pass
captured = capsys.readouterr()
assert "Starting phase_under_test component=component-1" in captured.out
assert "Finished phase_under_test in " in captured.out
assert "component=component-1" in captured.out
def test_modal_cloud_finds_memory_mapped_compile_cache_files(
modal_cloud_module: Any,
tmp_path: Path,
) -> None:
"""Identify native cache files visible only in process memory maps."""
volume_root = tmp_path / "compile-cache"
mapped_library = volume_root / "triton" / "cuda_utils.so"
mapped_library.parent.mkdir(parents=True)
mapped_library.touch()
proc_root = tmp_path / "proc"
process_root = proc_root / "191"
process_root.mkdir(parents=True)
(process_root / "maps").write_text(
f"7f000000-7f001000 r-xp 00000000 00:1e 144 {mapped_library}\n"
f"7f001000-7f002000 r--p 00001000 00:1e 144 {mapped_library}\n",
encoding="utf-8",
)
mapped_files = modal_cloud_module._mapped_process_files_under(
volume_root,
proc_root=proc_root,
)
assert mapped_files == ((191, str(mapped_library)),)
def test_invoke_mapped_remote_engine_async_runs_explicit_mapped_phase_items(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped remote execution should run the explicit mapped phase once per item in order."""
observed_calls: list[tuple[str, dict[str, Any]]] = []
progress_updates: list[dict[str, Any]] = []
def fake_execute_subgraph_prompt(
payload: dict[str, Any],
hydrated_inputs: dict[str, Any],
node_mapping: Any = None,
) -> tuple[str]:
assert payload["payload_kind"] == "subgraph"
assert payload["suppress_status_stream"] is True
observed_calls.append((str(payload["component_id"]), dict(hydrated_inputs)))
return (f"done:{hydrated_inputs['remote_input_0']}",)
monkeypatch.setattr(
mapped_execution_module,
"_execute_subgraph_prompt",
fake_execute_subgraph_prompt,
)
monkeypatch.setattr(
mapped_execution_module,
"_emit_local_modal_progress",
lambda **kwargs: progress_updates.append(kwargs),
)
payload = {
"payload_kind": "mapped_subgraph",
"component_id": "6",
"prompt_id": "prompt-1",
"mapped_input": {"proxy_input_name": "remote_input_0", "io_type": "STRING"},
"boundary_outputs": [
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"static_to_mapped_boundaries": [],
"static_phase": {
"component_node_ids": [],
"subgraph_prompt": {},
"boundary_inputs": [],
"boundary_outputs": [],
"execute_node_ids": [],
},
"mapped_phase": {
"component_node_ids": ["7"],
"subgraph_prompt": {
"7": {
"class_type": "RemoteStringEcho",
"inputs": {"text": ["remote_input_0", 0]},
}
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "STRING",
"targets": [{"node_id": "7", "input_name": "text"}],
}
],
"boundary_outputs": [
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"execute_node_ids": ["7"],
},
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_mapped_remote_engine_async(
payload,
serialization_module.serialize_node_inputs(
{"remote_input_0": ["a", "b", "c", "d"]}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["done:a", "done:b", "done:c", "done:d"],
)
assert observed_calls == [
("6::item:0", {"remote_input_0": "a"}),
("6::item:1", {"remote_input_0": "b"}),
("6::item:2", {"remote_input_0": "c"}),
("6::item:3", {"remote_input_0": "d"}),
]
assert progress_updates[0]["value"] == 0.0
assert progress_updates[0].get("lane_id") is None
assert progress_updates[-1]["value"] == 4.0
@pytest.mark.parametrize(
("module_fixture_name", "aggregate_function_name"),
[
("remote_modal_app_module", "_aggregate_mapped_outputs"),
("modal_cloud_module", "_aggregate_mapped_phase_outputs"),
],
)
def test_mapped_latent_aggregation_preserves_scheduler_items_when_shapes_differ(
request: Any,
module_fixture_name: str,
aggregate_function_name: str,
) -> None:
"""Both execution paths should return heterogeneous LATENT values as scheduler items."""
torch = pytest.importorskip("torch")
target_module = request.getfixturevalue(module_fixture_name)
aggregate_outputs = getattr(target_module, aggregate_function_name)
first_latent = {"samples": torch.zeros((1, 4, 32, 32), dtype=torch.float32)}
second_latent = {"samples": torch.zeros((1, 4, 35, 35), dtype=torch.float32)}
payload = {
"boundary_outputs": [
{
"io_type": "LATENT",
"is_list": False,
"mapped_output": True,
"scheduler_is_list": True,
}
]
}
result = aggregate_outputs([(first_latent,), (second_latent,)], payload)
assert len(result) == 1
assert len(result[0]) == 2
assert result[0][0] is first_latent
assert result[0][1] is second_latent
def test_invoke_mapped_remote_engine_async_splits_int_inputs_for_direct_targets(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped remote execution should itemize list(INT) inputs even when ModalMapInput stays local."""
observed_calls: list[tuple[str, dict[str, Any]]] = []
def fake_execute_subgraph_prompt(
payload: dict[str, Any],
hydrated_inputs: dict[str, Any],
node_mapping: Any = None,
) -> tuple[str]:
observed_calls.append((str(payload["component_id"]), dict(hydrated_inputs)))
return (f"seed:{hydrated_inputs['remote_input_0']}",)
monkeypatch.setattr(
mapped_execution_module,
"_execute_subgraph_prompt",
fake_execute_subgraph_prompt,
)
payload = {
"payload_kind": "mapped_subgraph",
"component_id": "12",
"prompt_id": "prompt-1",
"mapped_input": {"proxy_input_name": "remote_input_0", "io_type": "INT"},
"boundary_outputs": [
{
"proxy_output_name": "12_latent",
"node_id": "12",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"static_to_mapped_boundaries": [],
"static_phase": {
"component_node_ids": [],
"subgraph_prompt": {},
"boundary_inputs": [],
"boundary_outputs": [],
"execute_node_ids": [],
},
"mapped_phase": {
"component_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "RemoteSampler",
"inputs": {"seed": ["remote_input_0", 0]},
}
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
}
],
"boundary_outputs": [
{
"proxy_output_name": "12_latent",
"node_id": "12",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"execute_node_ids": ["12"],
},
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_mapped_remote_engine_async(
payload,
serialization_module.serialize_node_inputs(
{"remote_input_0": [10, 11, 12]}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
[
"seed:10",
"seed:11",
"seed:12",
],
)
assert observed_calls == [
("12::item:0", {"remote_input_0": 10}),
("12::item:1", {"remote_input_0": 11}),
("12::item:2", {"remote_input_0": 12}),
]
def test_invoke_mapped_remote_engine_async_executes_static_branch_once(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped remote execution should run the explicit static phase once and inject its bridge outputs."""
observed_execute_node_ids: list[tuple[str, tuple[str, ...]]] = []
def fake_execute_subgraph_prompt(
payload: dict[str, Any],
hydrated_inputs: dict[str, Any],
node_mapping: Any = None,
) -> tuple[str, ...]:
observed_execute_node_ids.append(
(
str(payload["component_id"]),
tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])),
)
)
if str(payload["component_id"]).endswith("::static"):
assert tuple(payload.get("execute_node_ids", [])) == ("1", "3")
assert [output["proxy_output_name"] for output in payload.get("boundary_outputs", [])] == [
"3_text",
"static_input_0",
]
return ("static-output", "shared-model")
assert tuple(payload.get("execute_node_ids", [])) == ("7",)
assert hydrated_inputs["static_input_0"] == "shared-model"
assert [output["proxy_output_name"] for output in payload.get("boundary_outputs", [])] == ["7_text"]
return (f"mapped:{hydrated_inputs['remote_input_1']}",)
monkeypatch.setattr(
mapped_execution_module,
"_execute_subgraph_prompt",
fake_execute_subgraph_prompt,
)
payload = {
"payload_kind": "mapped_subgraph",
"component_id": "1",
"prompt_id": "prompt-1",
"mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "1",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "7", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["1", "3"],
"subgraph_prompt": {
"1": {"class_type": "RemoteModel", "inputs": {}},
"3": {"class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "3_text",
"node_id": "3",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
},
{
"proxy_output_name": "static_input_0",
"node_id": "1",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
},
],
"execute_node_ids": ["1", "3"],
},
"mapped_phase": {
"component_node_ids": ["6", "7"],
"subgraph_prompt": {
"6": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
"7": {
"class_type": "RemoteSampler",
"inputs": {"model": ["static_input_0", 0], "latent": ["6", 0]},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_1",
"io_type": "STRING",
"targets": [{"node_id": "6", "input_name": "value"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "7", "input_name": "model"}],
},
],
"boundary_outputs": [
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"execute_node_ids": ["7"],
},
"boundary_outputs": [
{
"proxy_output_name": "3_text",
"node_id": "3",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": False,
},
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
},
],
"execute_node_ids": ["3", "7"],
"static_execute_node_ids": ["1", "3"],
"mapped_execute_node_ids": ["7"],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_mapped_remote_engine_async(
payload,
serialization_module.serialize_node_inputs(
{"remote_input_1": ["a", "b"]}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
"static-output",
["mapped:a", "mapped:b"],
)
assert observed_execute_node_ids[0] == ("1::static", ("1", "3"))
assert observed_execute_node_ids[1:] == [
("1::item:0", ("7",)),
("1::item:1", ("7",)),
]
def test_modal_cloud_execute_mapped_subgraph_payload_injects_static_bridges(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""The cloud runtime should execute the static phase once and feed its outputs into each mapped item."""
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
def fake_execute_subgraph_prompt(
payload: dict[str, Any],
hydrated_inputs: dict[str, Any],
custom_nodes_root: Any,
status_callback: Any = None,
cancellation_event: Any = None,
interrupt_store: Any = None,
interrupt_flag_key: Any = None,
) -> tuple[str, ...]:
observed_calls.append(
(
str(payload["component_id"]),
tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])),
dict(hydrated_inputs),
)
)
if str(payload["component_id"]).endswith("::static"):
return ("static-output", "shared-model")
return (f"mapped:{hydrated_inputs['remote_input_1']}:{hydrated_inputs['static_input_0']}",)
monkeypatch.setattr(
_cloud_mapped_execution_owner(),
"_execute_prompt_subgraph",
fake_execute_subgraph_prompt,
)
payload = {
"payload_kind": "mapped_subgraph",
"component_id": "cloud-1",
"prompt_id": "prompt-1",
"mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "1",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "7", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["1", "3"],
"subgraph_prompt": {
"1": {"class_type": "RemoteModel", "inputs": {}},
"3": {"class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "3_text",
"node_id": "3",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
},
{
"proxy_output_name": "static_input_0",
"node_id": "1",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
},
],
"execute_node_ids": ["1", "3"],
},
"mapped_phase": {
"component_node_ids": ["6", "7"],
"subgraph_prompt": {
"6": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
"7": {
"class_type": "RemoteSampler",
"inputs": {"model": ["static_input_0", 0], "latent": ["6", 0]},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_1",
"io_type": "STRING",
"targets": [{"node_id": "6", "input_name": "value"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "7", "input_name": "model"}],
},
],
"boundary_outputs": [
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"execute_node_ids": ["7"],
},
"boundary_outputs": [
{
"proxy_output_name": "3_text",
"node_id": "3",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": False,
},
{
"proxy_output_name": "7_text",
"node_id": "7",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
},
],
"extra_data": {"client_id": "client-1"},
}
outputs = modal_cloud_module._execute_mapped_subgraph_payload(
payload,
{"remote_input_1": ["a", "b"]},
None,
)
assert outputs == (
"static-output",
["mapped:a:shared-model", "mapped:b:shared-model"],
)
assert observed_calls == [
("cloud-1::static", ("1", "3"), {}),
("cloud-1::item:0", ("7",), {"static_input_0": "shared-model", "remote_input_1": "a"}),
("cloud-1::item:1", ("7",), {"static_input_0": "shared-model", "remote_input_1": "b"}),
]
def test_modal_cloud_execute_mapped_subgraph_payload_preserves_assigned_lane_id(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped cloud progress should keep the caller-assigned lane id instead of collapsing to `0`."""
observed_progress_events: list[dict[str, Any]] = []
def fake_execute_subgraph_prompt(
payload: dict[str, Any],
hydrated_inputs: dict[str, Any],
custom_nodes_root: Any,
status_callback: Any = None,
cancellation_event: Any = None,
interrupt_store: Any = None,
interrupt_flag_key: Any = None,
) -> tuple[str, ...]:
del custom_nodes_root, cancellation_event, interrupt_store, interrupt_flag_key
if status_callback is not None:
status_callback(
{
"event_type": "node_progress",
"node_id": "12",
"display_node_id": "12",
"real_node_id": "12",
"value": 3.0,
"max": 9.0,
}
)
return (f"mapped:{hydrated_inputs['remote_input_1']}",)
monkeypatch.setattr(
_cloud_mapped_execution_owner(),
"_execute_prompt_subgraph",
fake_execute_subgraph_prompt,
)
payload = {
"payload_kind": "mapped_subgraph",
"component_id": "cloud-2",
"prompt_id": "prompt-1",
"mapped_progress_lane_id": "3",
"mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
"static_to_mapped_boundaries": [],
"mapped_phase": {
"component_node_ids": ["12", "39"],
"subgraph_prompt": {
"39": {
"class_type": "RemoteSampler",
"inputs": {"latent": ["remote_input_1", 0]},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_1",
"io_type": "STRING",
"targets": [{"node_id": "39", "input_name": "latent"}],
}
],
"boundary_outputs": [
{
"proxy_output_name": "39_text",
"node_id": "39",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
}
],
"execute_node_ids": ["39"],
},
"boundary_outputs": [
{
"proxy_output_name": "39_text",
"node_id": "39",
"output_index": 0,
"io_type": "STRING",
"is_list": False,
"mapped_output": True,
},
],
}
outputs = modal_cloud_module._execute_mapped_subgraph_payload(
payload,
{"remote_input_1": ["a"]},
None,
status_callback=lambda event: observed_progress_events.append(dict(event)),
)
assert outputs == (["mapped:a"],)
assert any(
event.get("event_type") == "node_progress"
and event.get("real_node_id") == "12"
and event.get("lane_id") == "3"
for event in observed_progress_events
)
assert any(
event.get("event_type") == "node_progress"
and event.get("clear") is True
and event.get("lane_id") == "3"
for event in observed_progress_events
)
def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies(
api_intercept_module: Any,
modal_executor_module: Any,
settings_module: Any,
sync_engine_module: Any,
tmp_path: Path,
) -> None:
"""A mapped input should not collapse local feedback back into the same remote proxy."""
settings = settings_module.ModalSyncSettings(
app_name="app",
auto_deploy=True,
allow_ephemeral_fallback=False,
enable_memory_snapshot=True,
enable_gpu_memory_snapshot=False,
execution_mode="local",
sync_custom_nodes=False,
volume_name="volume",
route_path="/modal/queue_prompt",
marker_property="is_modal_remote",
local_storage_root=tmp_path / "storage",
remote_storage_root="/storage",
custom_nodes_archive_name="custom_nodes_bundle.zip",
comfyui_root=None,
custom_nodes_dir=tmp_path / "custom_nodes",
)
settings.custom_nodes_dir.mkdir()
sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
fake_nodes_module = type(
"FakeNodesModule",
(),
{
"NODE_CLASS_MAPPINGS": {
"RemoteModel": _FakeRewriteRemoteModelNode,
"RemoteSampler": _FakeRewriteRemoteSamplerNode,
"LatentSource": _FakeRewriteLatentSourceNode,
"ModalMapInput": _FakeRewriteModalMapInputNode,
"LocalFeedback": _FakeRewriteLocalFeedbackNode,
"LocalSink": _FakeRewriteLocalSinkNode,
},
"NODE_DISPLAY_NAME_MAPPINGS": {},
},
)()
workflow = {
"nodes": [
{"id": 1, "properties": {"is_modal_remote": True}},
{"id": 2, "properties": {"is_modal_remote": False}},
{"id": 3, "properties": {"is_modal_remote": True}},
{"id": 4, "properties": {"is_modal_remote": False}},
{"id": 5, "properties": {"is_modal_remote": False}},
{"id": 6, "properties": {"is_modal_remote": False}},
{"id": 7, "properties": {"is_modal_remote": True}},
{"id": 8, "properties": {"is_modal_remote": False}},
{"id": 9, "properties": {"is_modal_remote": False}},
]
}
prompt = {
"1": {
"class_type": "RemoteModel",
"inputs": {},
"_meta": {"title": "Shared Model"},
},
"2": {
"class_type": "LatentSource",
"inputs": {},
"_meta": {"title": "Single Latent"},
},
"3": {
"class_type": "RemoteSampler",
"inputs": {"model": ["1", 0], "latent": ["2", 0]},
"_meta": {"title": "First Sampler"},
},
"4": {
"class_type": "LocalFeedback",
"inputs": {"image": ["3", 0]},
"_meta": {"title": "Local Feedback"},
},
"5": {
"class_type": "LatentSource",
"inputs": {},
"_meta": {"title": "Batch Values"},
},
"6": {
"class_type": "ModalMapInput",
"inputs": {"value": ["5", 0]},
"_meta": {"title": "Map Input"},
},
"7": {
"class_type": "RemoteSampler",
"inputs": {"model": ["1", 0], "latent": ["6", 0], "prompt": ["4", 0]},
"_meta": {"title": "Mapped Sampler"},
},
"8": {
"class_type": "LocalSink",
"inputs": {"image": ["7", 0]},
"_meta": {"title": "Local Sink"},
},
"9": {
"class_type": "LocalSink",
"inputs": {"image": ["3", 0]},
"_meta": {"title": "Independent Local Preview"},
},
}
rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal(
prompt=prompt,
workflow=workflow,
sync_engine=sync_engine,
settings=settings,
nodes_module=fake_nodes_module,
)
assert api_intercept_module._find_prompt_dependency_cycles(rewritten_prompt) == []
assert summary.mapped_component_ids == ["7"]
assert summary.component_execution_stages == [["3"], ["7"]]
assert rewritten_prompt["4"]["inputs"]["image"] == ["3", 0]
assert rewritten_prompt["7"]["inputs"]["remote_input_2"] == ["4", 0]
assert len(summary.parallel_local_branch_node_ids) == 1
passthrough_node_id = summary.parallel_local_branch_node_ids[0]
passthrough_node = rewritten_prompt[passthrough_node_id]
assert passthrough_node["class_type"] == (
api_intercept_module.MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID
)
assert passthrough_node["inputs"]["value"] == ["3", 0]
dispatch_context = passthrough_node["inputs"]["dispatch_context"]
assert dispatch_context["component_ids"] == ["7"]
assert rewritten_prompt["9"]["inputs"]["image"] == [passthrough_node_id, 0]
first_payload = rewritten_prompt["3"]["inputs"]["original_node_data"]
mapped_payload = rewritten_prompt["7"]["inputs"]["original_node_data"]
assert first_payload["component_node_ids"] == ["1", "3"]
assert mapped_payload["component_node_ids"] == ["7"]
assert first_payload["remote_local_gap_pool"] is True
assert first_payload["keepalive_after_remote_component"] is True
assert "stop_local_gap_keepalive_before_remote_component" not in first_payload
assert mapped_payload["remote_local_gap_pool"] is True
assert "keepalive_after_remote_component" not in mapped_payload
assert mapped_payload["stop_local_gap_keepalive_before_remote_component"] is True
assert mapped_payload["parallel_local_dispatch_group_id"] == (
dispatch_context["dispatch_group_id"]
)
assert mapped_payload["signal_parallel_local_dispatch"] is True
rehydrated_mapped_payload = modal_executor_module._rehydrate_proxy_payload(
mapped_payload,
unique_id="7",
)
assert rehydrated_mapped_payload["signal_parallel_local_dispatch"] is True
assert {
boundary_input["proxy_input_name"]
for boundary_input in mapped_payload["boundary_inputs"]
} == {"phase_bridge_0", "remote_input_1", "remote_input_2"}
def test_invoke_implicitly_mapped_subgraph_async_zips_batched_boundary_inputs(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Ordinary remote subgraphs should fan out when multiple boundary inputs arrive batched."""
observed_inputs: list[dict[str, Any]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
assert payload["payload_kind"] == "subgraph"
assert payload["suppress_status_stream"] is True
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
observed_inputs.append(hydrated_inputs)
return serialization_module.serialize_node_outputs(
(f"{hydrated_inputs['remote_input_0']}:{hydrated_inputs['remote_input_1']}",)
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "12",
"prompt_id": "prompt-1",
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {
"latent_image": ["remote_input_0", 0],
"seed": ["remote_input_1", 0],
},
}
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "LATENT",
"targets": [{"node_id": "12", "input_name": "latent_image"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
],
"boundary_outputs": [{"node_id": "12", "io_type": "STRING", "is_list": False}],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": ["latent-a", "latent-b"],
"remote_input_1": [10, 11],
}
),
)
)
assert observed_inputs == [
{"remote_input_0": "latent-a", "remote_input_1": 10},
{"remote_input_0": "latent-b", "remote_input_1": 11},
]
assert serialization_module.deserialize_node_outputs(response) == (
["latent-a:10", "latent-b:11"],
)
def test_implicit_mapping_preserves_create_video_frame_sequence(
remote_modal_app_module: Any,
caplog: pytest.LogCaptureFixture,
) -> None:
"""CreateVideo should receive one complete IMAGE frame sequence."""
torch = pytest.importorskip("torch")
frame_sequence = torch.zeros((124, 8, 8, 3), dtype=torch.float32)
payload = {
"payload_kind": "subgraph",
"component_id": "105:91",
"subgraph_prompt": {
"105:91": {
"class_type": "CreateVideo",
"inputs": {"images": ["105:10", 0]},
}
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "IMAGE",
"targets": [{"node_id": "105:91", "input_name": "images"}],
}
],
}
caplog.set_level(logging.INFO)
split_inputs = remote_modal_app_module._split_batch_boundary_inputs(
payload,
{"remote_input_0": frame_sequence},
)
assert split_inputs is None
assert "target sockets consume the complete tensor batch" in caplog.text
assert "105:91.images" in caplog.text
def test_implicitly_mapped_subgraph_shared_model_keeps_unbatched_sampler_single_run(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""A shared MODEL with mixed batch-size INT seeds should run sampler 4 once and sampler 12 four times."""
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
if execute_node_ids == ("4",):
return serialization_module.serialize_node_outputs(("sampler-4",))
if execute_node_ids == ("12",):
return serialization_module.serialize_node_outputs((f"sampler-12:{hydrated_inputs['remote_input_0']}",))
raise AssertionError(f"Unexpected execute nodes for implicit mapped regression: {execute_node_ids!r}")
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["4", "12", "17"],
"execute_node_ids": ["4", "12"],
"subgraph_prompt": {
"17": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
"4": {"class_type": "KSampler", "inputs": {"model": ["17", 0], "seed": 0}},
"12": {
"class_type": "KSampler",
"inputs": {"model": ["17", 0], "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "4", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "4", "io_type": "STRING", "is_list": False},
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11, 12, 13],
"remote_input_1": [28],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
"sampler-4",
["sampler-12:10", "sampler-12:11", "sampler-12:12", "sampler-12:13"],
)
assert observed_calls == [
("17::static", ("4",), {"remote_input_1": [28]}),
("17::item:0", ("12",), {"remote_input_0": 10, "remote_input_1": [28]}),
("17::item:1", ("12",), {"remote_input_0": 11, "remote_input_1": [28]}),
("17::item:2", ("12",), {"remote_input_0": 12, "remote_input_1": [28]}),
("17::item:3", ("12",), {"remote_input_0": 13, "remote_input_1": [28]}),
]
def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped remote scheduling should seed one bound worker lane before sending per-item calls there."""
observed_calls: list[tuple[str, str, bool, tuple[str, ...], dict[str, Any]]] = []
observed_lane_setup_starts: list[tuple[str, int]] = []
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
monkeypatch.setattr(remote_modal_app_module, "modal", object())
monkeypatch.setattr(mapped_execution_module, "modal", object())
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
monkeypatch.setattr(
mapped_execution_module,
"_lookup_deployed_remote_engine",
lambda payload, affinity_key_override=None: (
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
),
)
async def fake_invoke_bound_remote_engine_async(
remote_engine: Any,
payload: dict[str, Any],
kwargs_payload: bytes,
) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append(
(
str(remote_engine),
str(payload["component_id"]),
bool(payload.get("clear_remote_session")),
execute_node_ids,
hydrated_inputs,
)
)
if str(payload["component_id"]).startswith("17::seed:"):
return serialization_module.serialize_node_outputs(())
if str(payload["component_id"]).startswith("17::cleanup:"):
return serialization_module.serialize_node_outputs(())
if execute_node_ids == ("12",):
await asyncio.sleep(0)
return serialization_module.serialize_node_outputs((f"{remote_engine}:{hydrated_inputs['remote_input_0']}",))
raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")
monkeypatch.setattr(
remote_modal_app_module,
"_invoke_bound_remote_engine_async",
fake_invoke_bound_remote_engine_async,
)
monkeypatch.setattr(
mapped_execution_module,
"_emit_local_mapped_lane_progress_start",
lambda payload, lane_index, item_index=None: observed_lane_setup_starts.append(
(str(payload["component_id"]), int(lane_index))
),
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {"model": 0, "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "12", "input_name": "model"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "12", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["4"],
"subgraph_prompt": {
"4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"session_output": True,
"preview_target_node_ids": [],
}
],
"execute_node_ids": ["4"],
},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11, 12, 13],
"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
}
),
)
)
seed_calls = [call for call in observed_calls if call[1].startswith("17::seed:")]
item_calls = [call for call in observed_calls if call[1].startswith("17::item:")]
cleanup_calls = [call for call in observed_calls if call[1].startswith("17::cleanup:")]
assert sorted(seed_calls) == [
(
"engine:worker-pool:slot:0",
"17::seed:0",
False,
("4",),
{"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}},
),
(
"engine:worker-pool:slot:1",
"17::seed:1",
False,
("4",),
{"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}},
),
]
assert {call[0] for call in item_calls} <= {"engine:worker-pool:slot:0", "engine:worker-pool:slot:1"}
assert {call[0] for call in item_calls} <= {call[0] for call in seed_calls}
assert sorted(cleanup_calls) == [
("engine:worker-pool:slot:0", "17::cleanup:0", True, (), {}),
("engine:worker-pool:slot:1", "17::cleanup:1", True, (), {}),
]
assert sorted(observed_lane_setup_starts) == [("17", 0), ("17", 1)]
response_outputs = serialization_module.deserialize_node_outputs(response)[0]
assert len(response_outputs) == 4
assert [output.rsplit(":", 1)[-1] for output in response_outputs] == ["10", "11", "12", "13"]
def test_implicitly_mapped_subgraph_allows_ready_lanes_to_start_before_slowest_seed_finishes(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""A fast seeded lane should start item execution before another lane finishes seeding."""
observed_events: list[tuple[str, str, str]] = []
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
monkeypatch.setattr(remote_modal_app_module, "modal", object())
monkeypatch.setattr(mapped_execution_module, "modal", object())
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
monkeypatch.setattr(
mapped_execution_module,
"_lookup_deployed_remote_engine",
lambda payload, affinity_key_override=None: (
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
),
)
async def fake_invoke_bound_remote_engine_async(
remote_engine: Any,
payload: dict[str, Any],
kwargs_payload: bytes,
) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
component_id = str(payload["component_id"])
if component_id == "17::seed:0":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(0.05)
observed_events.append(("seed_done", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id == "17::seed:1":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(0)
observed_events.append(("seed_done", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::cleanup:"):
observed_events.append(("cleanup", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::item:"):
observed_events.append(("item", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(
(f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
)
raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")
monkeypatch.setattr(
remote_modal_app_module,
"_invoke_bound_remote_engine_async",
fake_invoke_bound_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {"model": 0, "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "12", "input_name": "model"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "12", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["4"],
"subgraph_prompt": {
"4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"session_output": True,
"preview_target_node_ids": [],
}
],
"execute_node_ids": ["4"],
},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11, 12, 13],
"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
}
),
)
)
assert len(serialization_module.deserialize_node_outputs(response)[0]) == 4
assert ("seed_done", "17::seed:1", "engine:worker-pool:slot:1") in observed_events
first_item_index = next(
index for index, event in enumerate(observed_events) if event[0] == "item"
)
fast_seed_done_index = observed_events.index(
("seed_done", "17::seed:1", "engine:worker-pool:slot:1")
)
assert first_item_index > fast_seed_done_index
def test_implicitly_mapped_subgraph_returns_once_all_items_finish_without_waiting_for_unused_lane(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""Prompt completion should not wait for a slow seeded lane after other lanes finish all mapped items."""
observed_events: list[tuple[str, str, str]] = []
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
monkeypatch.setattr(remote_modal_app_module, "modal", object())
monkeypatch.setattr(mapped_execution_module, "modal", object())
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
monkeypatch.setattr(
mapped_execution_module,
"_lookup_deployed_remote_engine",
lambda payload, affinity_key_override=None: (
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
),
)
async def fake_invoke_bound_remote_engine_async(
remote_engine: Any,
payload: dict[str, Any],
kwargs_payload: bytes,
) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
component_id = str(payload["component_id"])
if component_id == "17::seed:0":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(1.0)
observed_events.append(("seed_done", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id == "17::seed:1":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(0)
observed_events.append(("seed_done", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::cleanup:"):
observed_events.append(("cleanup", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::item:"):
observed_events.append(("item", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(
(f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
)
raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")
monkeypatch.setattr(
remote_modal_app_module,
"_invoke_bound_remote_engine_async",
fake_invoke_bound_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {"model": 0, "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "12", "input_name": "model"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "12", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["4"],
"subgraph_prompt": {
"4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"session_output": True,
"preview_target_node_ids": [],
}
],
"execute_node_ids": ["4"],
},
}
started_at = time.perf_counter()
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11],
"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
}
),
)
)
elapsed = time.perf_counter() - started_at
assert serialization_module.deserialize_node_outputs(response)[0] == [
"engine:worker-pool:slot:1:10",
"engine:worker-pool:slot:1:11",
]
assert elapsed < 0.5
assert ("item", "17::item:0", "engine:worker-pool:slot:1") in observed_events
assert ("cleanup", "17::cleanup:1", "engine:worker-pool:slot:1") in observed_events
assert ("cleanup", "17::cleanup:0", "engine:worker-pool:slot:0") not in observed_events
assert ("seed_done", "17::seed:0", "engine:worker-pool:slot:0") not in observed_events
def test_implicitly_mapped_subgraph_ignores_late_lane_failure_after_all_items_complete(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""A detached late lane must not fail the prompt after all mapped outputs are already complete."""
observed_events: list[tuple[str, str, str]] = []
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
monkeypatch.setattr(remote_modal_app_module, "modal", object())
monkeypatch.setattr(mapped_execution_module, "modal", object())
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
monkeypatch.setattr(
mapped_execution_module,
"_lookup_deployed_remote_engine",
lambda payload, affinity_key_override=None: (
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
),
)
async def fake_invoke_bound_remote_engine_async(
remote_engine: Any,
payload: dict[str, Any],
kwargs_payload: bytes,
) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
component_id = str(payload["component_id"])
if component_id == "17::seed:0":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(0.05)
raise RuntimeError("late lane failed after prompt completion")
if component_id == "17::seed:1":
observed_events.append(("seed_start", component_id, str(remote_engine)))
await asyncio.sleep(0)
observed_events.append(("seed_done", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::cleanup:"):
observed_events.append(("cleanup", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(())
if component_id.startswith("17::item:"):
observed_events.append(("item", component_id, str(remote_engine)))
return serialization_module.serialize_node_outputs(
(f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
)
raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")
monkeypatch.setattr(
remote_modal_app_module,
"_invoke_bound_remote_engine_async",
fake_invoke_bound_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {"model": 0, "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "static_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "12", "input_name": "model"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
"static_to_mapped_boundaries": [
{
"proxy_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"targets": [{"node_id": "12", "input_name": "model"}],
}
],
"static_phase": {
"component_node_ids": ["4"],
"subgraph_prompt": {
"4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
},
"boundary_inputs": [],
"boundary_outputs": [
{
"proxy_output_name": "static_input_0",
"node_id": "4",
"output_index": 0,
"io_type": "MODEL",
"is_list": False,
"session_output": True,
"preview_target_node_ids": [],
}
],
"execute_node_ids": ["4"],
},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11],
"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
}
),
)
)
assert serialization_module.deserialize_node_outputs(response)[0] == [
"engine:worker-pool:slot:1:10",
"engine:worker-pool:slot:1:11",
]
assert ("cleanup", "17::cleanup:1", "engine:worker-pool:slot:1") in observed_events
assert ("cleanup", "17::cleanup:0", "engine:worker-pool:slot:0") not in observed_events
def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Implicit fan-out should not split components whose list boundary lands on INPUT_IS_LIST nodes."""
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
fake_nodes_module = type(
"FakeNodesModule",
(),
{
"NODE_CLASS_MAPPINGS": {
"ImplicitBatchListSource": _ImplicitBatchListSourceNode,
"ImplicitBatchScalarConsumer": _ImplicitBatchScalarConsumerNode,
"ImplicitBatchListConsumer": _ImplicitBatchListConsumerNode,
}
},
)()
monkeypatch.setattr(mapped_execution_module, "_load_nodes_module", lambda: fake_nodes_module)
async def fake_invoke_remote_engine_async(
payload: dict[str, Any],
kwargs_payload: bytes,
*,
allow_implicit_mapping: bool = True,
) -> bytes:
del allow_implicit_mapping
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
return serialization_module.serialize_node_outputs(
("scalar-output", ["list-output:0", "list-output:1", "list-output:2"])
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["4", "12", "17"],
"execute_node_ids": ["4", "12"],
"subgraph_prompt": {
"17": {"class_type": "ImplicitBatchListSource", "inputs": {"values": 0}},
"4": {
"class_type": "ImplicitBatchScalarConsumer",
"inputs": {"value": ["17", 0]},
},
"12": {
"class_type": "ImplicitBatchListConsumer",
"inputs": {"values": ["17", 1]},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "17", "input_name": "values"}],
}
],
"boundary_outputs": [
{"node_id": "4", "io_type": "STRING", "is_list": False},
{"node_id": "12", "io_type": "STRING", "is_list": True},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
"scalar-output",
["list-output:0", "list-output:1", "list-output:2"],
)
assert observed_calls == [
("17", ("4", "12"), {"remote_input_0": [10, 11, 12]}),
]
def test_invoke_remote_engine_async_bypasses_implicit_mapping_once_for_input_is_list_targets(
remote_modal_app_module: Any,
mapped_execution_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""The async dispatcher should not recurse when implicit fan-out is suppressed."""
observed_calls: list[tuple[bool, dict[str, Any]]] = []
fake_nodes_module = type(
"FakeNodesModule",
(),
{
"NODE_CLASS_MAPPINGS": {
"ImplicitBatchListSource": _ImplicitBatchListSourceNode,
"ImplicitBatchScalarConsumer": _ImplicitBatchScalarConsumerNode,
"ImplicitBatchListConsumer": _ImplicitBatchListConsumerNode,
}
},
)()
monkeypatch.setattr(mapped_execution_module, "_load_nodes_module", lambda: fake_nodes_module)
def fake_invoke_remote_engine(
payload: dict[str, Any],
kwargs_payload: bytes,
*,
allow_implicit_mapping: bool = True,
) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
observed_calls.append((allow_implicit_mapping, hydrated_inputs))
return serialization_module.serialize_node_outputs(("ordinary-subgraph",))
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine",
fake_invoke_remote_engine,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["4", "12", "17"],
"execute_node_ids": ["4", "12"],
"subgraph_prompt": {
"17": {"class_type": "ImplicitBatchListSource", "inputs": {"values": 0}},
"4": {
"class_type": "ImplicitBatchScalarConsumer",
"inputs": {"value": ["17", 0]},
},
"12": {
"class_type": "ImplicitBatchListConsumer",
"inputs": {"values": ["17", 1]},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "17", "input_name": "values"}],
}
],
"boundary_outputs": [
{"node_id": "4", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module.invoke_remote_engine_async(
payload,
serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
)
)
assert serialization_module.deserialize_node_outputs(response) == ("ordinary-subgraph",)
assert observed_calls == [
(False, {"remote_input_0": [10, 11, 12]}),
]
def test_implicitly_mapped_subgraph_clears_remote_session_once_after_all_items_finish(
remote_modal_app_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""Implicit mapped execution should reserve remote-session cleanup for one final cleanup payload."""
observed_calls: list[tuple[str, bool, tuple[str, ...], dict[str, Any]]] = []
session_cleared = False
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
nonlocal session_cleared
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
clear_remote_session = bool(payload.get("clear_remote_session"))
observed_calls.append(
(str(payload["component_id"]), clear_remote_session, execute_node_ids, hydrated_inputs)
)
if clear_remote_session:
session_cleared = True
elif session_cleared:
raise session_state_module.RemoteSessionStateError(
"Remote session 'session-1' was not found."
)
if execute_node_ids == ("4",):
return serialization_module.serialize_node_outputs(("sampler-4",))
if execute_node_ids == ("12",):
return serialization_module.serialize_node_outputs((f"sampler-12:{hydrated_inputs['remote_input_0']}",))
if execute_node_ids == ():
return serialization_module.serialize_node_outputs(())
raise AssertionError(f"Unexpected execute nodes for implicit cleanup regression: {execute_node_ids!r}")
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["4", "12", "17"],
"execute_node_ids": ["4", "12"],
"subgraph_prompt": {
"17": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
"4": {"class_type": "KSampler", "inputs": {"model": ["17", 0], "seed": 0}},
"12": {
"class_type": "KSampler",
"inputs": {"model": ["17", 0], "seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "4", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "4", "io_type": "STRING", "is_list": False},
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [10, 11, 12],
"remote_input_1": [28],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
"sampler-4",
["sampler-12:10", "sampler-12:11", "sampler-12:12"],
)
assert [
(component_id, execute_node_ids)
for component_id, _, execute_node_ids, _ in observed_calls
] == [
("17::static", ("4",)),
("17::item:0", ("12",)),
("17::item:1", ("12",)),
("17::item:2", ("12",)),
("17::cleanup", ()),
]
assert all(
not clear_remote_session
for component_id, clear_remote_session, _, _ in observed_calls
if component_id != "17::cleanup"
)
assert observed_calls[-1] == ("17::cleanup", True, (), {})
def test_implicitly_mapped_subgraph_stops_queued_items_after_local_interrupt(
remote_modal_app_module: Any,
serialization_module: Any,
session_state_module: Any,
monkeypatch: Any,
) -> None:
"""Implicit mapped cancellation should not dispatch queued items or cleanup work."""
class FakeInterrupt(Exception):
"""Stand-in for ComfyUI's InterruptProcessingException."""
observed_calls: list[str] = []
local_interrupted = False
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
nonlocal local_interrupted
del kwargs_payload
component_id = str(payload["component_id"])
observed_calls.append(component_id)
if component_id.endswith("::cleanup"):
raise AssertionError("interrupted mapped execution should not run cleanup")
local_interrupted = True
return serialization_module.serialize_node_outputs((component_id,))
def fake_local_processing_interrupted() -> bool:
return local_interrupted
monkeypatch.setattr(
remote_modal_app_module,
"boost_mapped_component_warmup",
lambda payload, *, total_items, reason: (1, 1),
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
monkeypatch.setattr(
remote_modal_app_module,
"_local_processing_interrupted",
fake_local_processing_interrupted,
)
monkeypatch.setattr(
remote_modal_app_module,
"_raise_local_interrupt",
lambda: (_ for _ in ()).throw(FakeInterrupt()),
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {"seed": 0},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
],
"boundary_outputs": [{"node_id": "12", "io_type": "STRING", "is_list": False}],
"extra_data": {"client_id": "client-1"},
"remote_session": session_state_module.RemoteSessionHandle(
session_id="session-1",
prompt_id="prompt-1",
owner_component_id="17",
).to_payload(),
"clear_remote_session": True,
}
with pytest.raises(FakeInterrupt):
asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
)
)
assert observed_calls == ["17::item:0"]
def test_implicitly_mapped_subgraph_keeps_conditioning_lists_broadcast(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Implicit mapped execution must not split list-backed CONDITIONING inputs per item."""
conditioning = [
["cond-a", {"pooled_output": "pool-a"}],
["cond-b", {"pooled_output": "pool-b"}],
]
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
if execute_node_ids != ("12",):
raise AssertionError(
f"Unexpected execute nodes for conditioning implicit mapped regression: {execute_node_ids!r}"
)
return serialization_module.serialize_node_outputs(
(f"{hydrated_inputs['remote_input_1']}:{len(hydrated_inputs['remote_input_0'])}",)
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {
"positive": 0,
"seed": 0,
},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "CONDITIONING",
"targets": [{"node_id": "12", "input_name": "positive"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": conditioning,
"remote_input_1": [10, 11, 12, 13],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["10:2", "11:2", "12:2", "13:2"],
)
assert observed_calls == [
("17::item:0", ("12",), {"remote_input_0": conditioning, "remote_input_1": 10}),
("17::item:1", ("12",), {"remote_input_0": conditioning, "remote_input_1": 11}),
("17::item:2", ("12",), {"remote_input_0": conditioning, "remote_input_1": 12}),
("17::item:3", ("12",), {"remote_input_0": conditioning, "remote_input_1": 13}),
]
def test_implicitly_mapped_subgraph_splits_mapped_conditioning_outputs(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Mapped CONDITIONING aggregates should feed downstream mapped items in lockstep."""
conditioning_items = [
[["cond-a", {"pooled_output": "pool-a"}]],
[["cond-b", {"pooled_output": "pool-b"}]],
[["cond-c", {"pooled_output": "pool-c"}]],
]
mapped_conditioning = serialization_module.join_mapped_values(
conditioning_items,
"CONDITIONING",
is_list=False,
)
latent_items = ["latent-a", "latent-b", "latent-c"]
observed_calls: list[tuple[str, dict[str, Any]]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
observed_calls.append((str(payload["component_id"]), hydrated_inputs))
return serialization_module.serialize_node_outputs(
(f"{hydrated_inputs['remote_input_0'][0][0]}:{hydrated_inputs['remote_input_1']}",)
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {
"positive": 0,
"latent_image": 0,
},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "CONDITIONING",
"targets": [{"node_id": "12", "input_name": "positive"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "LATENT",
"targets": [{"node_id": "12", "input_name": "latent_image"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": mapped_conditioning,
"remote_input_1": latent_items,
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["cond-a:latent-a", "cond-b:latent-b", "cond-c:latent-c"],
)
assert observed_calls == [
(
"17::item:0",
{"remote_input_0": conditioning_items[0], "remote_input_1": latent_items[0]},
),
(
"17::item:1",
{"remote_input_0": conditioning_items[1], "remote_input_1": latent_items[1]},
),
(
"17::item:2",
{"remote_input_0": conditioning_items[2], "remote_input_1": latent_items[2]},
),
]
def test_implicitly_mapped_subgraph_splits_wildcard_latent_lists_for_latent_targets(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Wildcard boundary lists of LATENT-like mappings should still itemize for LATENT sockets."""
latent_items = [
{"samples": "latent-a", "batch_index": [0]},
{"samples": "latent-b", "batch_index": [1]},
{"samples": "latent-c", "batch_index": [2]},
]
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
fake_nodes_module = type(
"FakeNodesModule",
(),
{"NODE_CLASS_MAPPINGS": {"KSampler": _FakeImplicitBatchKSamplerNode}},
)()
monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
if execute_node_ids != ("12",):
raise AssertionError(
f"Unexpected execute nodes for wildcard LATENT implicit mapped regression: {execute_node_ids!r}"
)
return serialization_module.serialize_node_outputs(
(f"{hydrated_inputs['remote_input_1']}:{hydrated_inputs['remote_input_0']['samples']}",)
)
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {
"latent_image": 0,
"seed": 0,
},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "*",
"targets": [{"node_id": "12", "input_name": "latent_image"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": latent_items,
"remote_input_1": [10, 11, 12],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["10:latent-a", "11:latent-b", "12:latent-c"],
)
assert observed_calls == [
("17::item:0", ("12",), {"remote_input_0": latent_items[0], "remote_input_1": 10}),
("17::item:1", ("12",), {"remote_input_0": latent_items[1], "remote_input_1": 11}),
("17::item:2", ("12",), {"remote_input_0": latent_items[2], "remote_input_1": 12}),
]
def test_implicitly_mapped_subgraph_splits_session_ref_lists_for_nontransportable_inputs(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Implicit mapped execution should itemize lists of remote session refs even for MODEL/CONDITIONING."""
model_ref = {
"__comfy_modal_remote_session_value_ref__": True,
"session_id": "session-1",
"node_id": "31",
"output_index": 0,
}
conditioning_ref = {
"__comfy_modal_remote_session_value_ref__": True,
"session_id": "session-1",
"node_id": "2",
"output_index": 0,
}
observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
if execute_node_ids != ("12",):
raise AssertionError(
f"Unexpected execute nodes for implicit session-ref regression: {execute_node_ids!r}"
)
return serialization_module.serialize_node_outputs((f"seed:{hydrated_inputs['remote_input_2']}",))
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "17",
"prompt_id": "prompt-1",
"component_node_ids": ["12"],
"execute_node_ids": ["12"],
"subgraph_prompt": {
"12": {
"class_type": "KSampler",
"inputs": {
"model": 0,
"positive": 0,
"seed": 0,
},
},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_0",
"io_type": "MODEL",
"targets": [{"node_id": "12", "input_name": "model"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "CONDITIONING",
"targets": [{"node_id": "12", "input_name": "positive"}],
},
{
"proxy_input_name": "remote_input_2",
"io_type": "INT",
"targets": [{"node_id": "12", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "12", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"remote_input_0": [model_ref, model_ref, model_ref],
"remote_input_1": [conditioning_ref, conditioning_ref, conditioning_ref],
"remote_input_2": [10, 11, 12],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["seed:10", "seed:11", "seed:12"],
)
assert observed_calls == [
(
"17::item:0",
("12",),
{"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 10},
),
(
"17::item:1",
("12",),
{"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 11},
),
(
"17::item:2",
("12",),
{"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 12},
),
]
def test_implicitly_mapped_subgraph_splits_singleton_wrapped_bridge_ref_lists(
remote_modal_app_module: Any,
serialization_module: Any,
monkeypatch: Any,
) -> None:
"""Comfy output wrappers around bridge-ref lists should still itemize mapped conditioning."""
first_conditioning_ref = {
"__comfy_modal_remote_session_bridge_ref__": True,
"bridge_key": "RSB_first",
"session_id": "session-1",
"node_id": "508",
"output_index": 0,
}
second_conditioning_ref = {
"__comfy_modal_remote_session_bridge_ref__": True,
"bridge_key": "RSB_second",
"session_id": "session-1",
"node_id": "508",
"output_index": 0,
}
observed_calls: list[tuple[str, dict[str, Any]]] = []
async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
observed_calls.append((str(payload["component_id"]), hydrated_inputs))
return serialization_module.serialize_node_outputs((f"seed:{hydrated_inputs['remote_input_1']}",))
monkeypatch.setattr(
remote_modal_app_module,
"invoke_remote_engine_async",
fake_invoke_remote_engine_async,
)
payload = {
"payload_kind": "subgraph",
"component_id": "507",
"prompt_id": "prompt-1",
"component_node_ids": ["507"],
"execute_node_ids": ["507"],
"subgraph_prompt": {
"507": {
"class_type": "KSampler",
"inputs": {
"positive": 0,
"seed": 0,
},
},
},
"boundary_inputs": [
{
"proxy_input_name": "phase_bridge_3",
"io_type": "CONDITIONING",
"targets": [{"node_id": "507", "input_name": "positive"}],
},
{
"proxy_input_name": "remote_input_1",
"io_type": "INT",
"targets": [{"node_id": "507", "input_name": "seed"}],
},
],
"boundary_outputs": [
{"node_id": "507", "io_type": "STRING", "is_list": False},
],
"extra_data": {"client_id": "client-1"},
}
response = asyncio.run(
remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
payload,
serialization_module.serialize_node_inputs(
{
"phase_bridge_3": [[first_conditioning_ref, second_conditioning_ref]],
"remote_input_1": [10, 11],
}
),
)
)
assert serialization_module.deserialize_node_outputs(response) == (
["seed:10", "seed:11"],
)
assert observed_calls == [
(
"507::item:0",
{"phase_bridge_3": first_conditioning_ref, "remote_input_1": 10},
),
(
"507::item:1",
{"phase_bridge_3": second_conditioning_ref, "remote_input_1": 11},
),
]
@pytest.mark.parametrize(
("module_fixture_name",),
[
("local_execution_module",),
("modal_cloud_module",),
],
)
def test_trim_subgraph_payload_to_required_nodes_drops_unrelated_mapped_branch(
request: Any,
module_fixture_name: str,
) -> None:
"""Static or per-item sub-runs should exclude unrelated nodes from the mapped sibling branch."""
target_module = request.getfixturevalue(module_fixture_name)
payload = {
"component_id": "1::static",
"component_node_ids": ["1", "2", "3", "7"],
"subgraph_prompt": {
"1": {"class_type": "LoadDiffusionModel", "inputs": {}},
"2": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
"3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "steps": 20}},
"7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "latent_image": ["2", 0]}},
},
"boundary_inputs": [
{
"proxy_input_name": "remote_input_1",
"targets": [{"node_id": "2", "input_name": "value"}],
}
],
"boundary_outputs": [
{"node_id": "3", "output_index": 0, "io_type": "LATENT", "is_list": False},
],
"execute_node_ids": ["3"],
"mapped_execute_node_ids": ["7"],
"static_execute_node_ids": ["3"],
}
trimmed_payload = target_module._trim_subgraph_payload_to_required_nodes(payload)
assert trimmed_payload["component_node_ids"] == ["1", "3"]
assert list(trimmed_payload["subgraph_prompt"].keys()) == ["1", "3"]
assert trimmed_payload["boundary_inputs"] == []
assert trimmed_payload["boundary_outputs"] == payload["boundary_outputs"]
assert trimmed_payload["execute_node_ids"] == ["3"]
assert trimmed_payload["mapped_execute_node_ids"] == []
assert trimmed_payload["static_execute_node_ids"] == ["3"]
Reference in New Issue View Git Blame Copy Permalink