fix: split mapped components with local reentry
This commit is contained in:
Ken Simpson
+23
-2
@@ -2672,8 +2672,13 @@ def _build_component_payload(
|
||||
return payload
|
||||
|
||||
def build_phase_payloads_for_transportable_splits() -> list[dict[str, Any]] | None:
|
||||
"""Return ordered split-proxy phase payloads for a non-mapped coarse component."""
|
||||
if component.mapped_boundary_input_name is not None or len(component.execute_node_ids) <= 1:
|
||||
"""Return ordered split-proxy phase payloads for components with local feedback."""
|
||||
if len(component.execute_node_ids) <= 1:
|
||||
return None
|
||||
if component.mapped_boundary_input_name is not None and not _component_has_local_reentry_dependency(
|
||||
prompt=signature_prompt,
|
||||
component=component,
|
||||
):
|
||||
return None
|
||||
if component.local_tap_node_ids:
|
||||
if not _component_has_local_reentry_dependency(
|
||||
@@ -3320,6 +3325,19 @@ def _rewrite_component_into_proxy(
|
||||
is_output_node=contains_output_node(list(phase_payload["component_node_ids"])),
|
||||
)
|
||||
|
||||
if component.mapped_boundary_source_node_id is not None:
|
||||
mapped_node_id_set = set(component.mapped_node_ids)
|
||||
for phase_payload in phase_payloads:
|
||||
phase_node_ids = {str(node_id) for node_id in phase_payload.get("component_node_ids", [])}
|
||||
if not (phase_node_ids & mapped_node_id_set):
|
||||
continue
|
||||
register_modal_map_input_warmup_context(
|
||||
component.mapped_boundary_source_node_id,
|
||||
phase_payload,
|
||||
str(component.mapped_boundary_input_io_type or "*"),
|
||||
)
|
||||
break
|
||||
|
||||
for node_id, prompt_node in list(rewritten_prompt.items()):
|
||||
if node_id in component_proxy_node_ids:
|
||||
continue
|
||||
@@ -3611,12 +3629,15 @@ def rewrite_prompt_for_modal(
|
||||
continue
|
||||
if isinstance(split_proxy_payloads, list):
|
||||
summary.remote_component_ids.extend(proxy_node_ids)
|
||||
mapped_node_id_set = set(component.mapped_node_ids)
|
||||
for phase_payload in split_proxy_payloads:
|
||||
phase_proxy_node_id = str(phase_payload["component_id"])
|
||||
phase_component_node_ids = [str(node_id) for node_id in phase_payload["component_node_ids"]]
|
||||
summary.component_node_ids_by_representative[phase_proxy_node_id] = phase_component_node_ids
|
||||
for node_id in phase_component_node_ids:
|
||||
summary.rewritten_node_id_map[node_id] = phase_proxy_node_id
|
||||
if mapped_node_id_set and mapped_node_id_set.intersection(phase_component_node_ids):
|
||||
mapped_proxy_component_ids.add(phase_proxy_node_id)
|
||||
continue
|
||||
|
||||
summary.remote_component_ids.extend(proxy_node_ids)
|
||||
|
||||
@@ -272,3 +272,4 @@
|
||||
- Cancellation can also arrive before a Modal payload reaches the registered remote-call wrapper, especially during first-run auto-deploy or deployed-app readiness polling. The outer local waiter must write the shared cancel flag and enforce the same local release grace during deploy/provisioning, and the worker must refuse to start the remote payload if cancellation was already observed.
|
||||
- Users need immediate cancellation feedback even when Modal cleanup is still cooperative. The frontend should intercept ComfyUI `/interrupt` requests, mark the prompt as `cancelling`, and ignore nonterminal late progress until the native terminal event clears the UI.
|
||||
- Prompt-level Modal cancellation has to release both the remote work and the local prompt admission gate. If a user cancels during Modal deploy/readiness, ComfyUI can mark the prompt interrupted while the background Modal call still owns the local gate, leaving the next prompt stuck behind an already-cancelled prompt id.
|
||||
- `ModalMapInput` does not remove the need for ordered split proxies. If a mapped remote component has a local re-entry dependency, such as an early sampler preview feeding a local text/image node that returns into a later mapped sampler, the static side must still split into ordered phases before the mapped phase or the rewritten prompt can create a false dependency cycle.
|
||||
|
||||
@@ -155,6 +155,14 @@ class _FakeRewriteLocalSinkNode:
|
||||
OUTPUT_IS_LIST = (False,)
|
||||
|
||||
|
||||
class _FakeRewriteLocalFeedbackNode:
|
||||
"""Fake local node that turns a remote latent into a transportable remote input."""
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
OUTPUT_IS_LIST = (False,)
|
||||
|
||||
|
||||
def test_modal_proxy_waits_for_active_prompt_before_starting_next_prompt(
|
||||
modal_executor_module: Any,
|
||||
) -> None:
|
||||
@@ -7555,6 +7563,125 @@ def test_split_hybrid_proxies_allow_local_downstream_work_before_mapped_completi
|
||||
]
|
||||
|
||||
|
||||
def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies(
|
||||
api_intercept_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}},
|
||||
]
|
||||
}
|
||||
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"},
|
||||
},
|
||||
}
|
||||
|
||||
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]
|
||||
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 {
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user