fix: detect modal markers inside subgraphs
This commit is contained in:
@@ -8,7 +8,7 @@
|
||||
## What is implemented
|
||||
|
||||
- A frontend extension in [web/modal_toggle.js](/home/user/git/ComfyUI-Modal/web/modal_toggle.js) that injects a `Run on Modal` toggle into every non-internal node, draws live remote-state overlays on marked nodes, and submits queued prompts to `/modal/queue_prompt`.
|
||||
- A backend route in [api_intercept.py](/home/user/git/ComfyUI-Modal/api_intercept.py) that reads the workflow snapshot from `extra_pnginfo.workflow`, identifies connected remote-marked node components, syncs referenced assets, and rewrites each component into one signature-preserving Modal proxy node.
|
||||
- A backend route in [api_intercept.py](/home/user/git/ComfyUI-Modal/api_intercept.py) that reads the workflow snapshot from `extra_pnginfo.workflow`, recursively discovers Modal markers even inside nested subgraph workflow metadata, identifies connected remote-marked node components, syncs referenced assets, and rewrites each component into one signature-preserving Modal proxy node.
|
||||
- A dynamic proxy registry in [modal_executor_node.py](/home/user/git/ComfyUI-Modal/modal_executor_node.py) that mirrors the exported output count and output types of each remote component so ComfyUI validation still succeeds after rewrite.
|
||||
- A content-addressable sync engine in [sync_engine.py](/home/user/git/ComfyUI-Modal/sync_engine.py) that mirrors model files and a zipped `custom_nodes/` bundle into storage keyed by SHA256.
|
||||
- A remote runtime skeleton in [remote/modal_app.py](/home/user/git/ComfyUI-Modal/remote/modal_app.py) that supports local fallback execution for tests and can call into Modal when the SDK is available.
|
||||
@@ -35,7 +35,7 @@
|
||||
## Runtime behavior
|
||||
|
||||
1. Queueing a prompt posts to `/modal/queue_prompt`.
|
||||
2. The route loads the workflow metadata from `extra_data.extra_pnginfo.workflow`.
|
||||
2. The route loads the workflow metadata from `extra_data.extra_pnginfo.workflow`, including nested saved subgraph fragments.
|
||||
3. Connected regions of nodes with `properties.is_modal_remote = true` are replaced with generated `ModalUniversalExecutor_<hash>` proxy nodes.
|
||||
4. File-like widget inputs ending in `.safetensors`, `.ckpt`, `.pt`, or `.vae` are resolved to local files and mirrored into storage.
|
||||
5. The executor serializes the component boundary inputs with `safetensors`, forwards the remote component payload to the runtime, and deserializes the exported boundary outputs back into ComfyUI values.
|
||||
@@ -171,6 +171,7 @@ Use the normal queue button in ComfyUI.
|
||||
The frontend extension intercepts queue submission and sends the prompt to `/modal/queue_prompt` instead of `/prompt`. On the backend:
|
||||
|
||||
1. The request handler inspects the workflow snapshot in `extra_pnginfo.workflow`.
|
||||
Modal marker discovery now walks nested saved subgraph fragments too, so nodes marked inside ComfyUI subgraphs still participate in remote rewrite.
|
||||
2. Connected remote-marked nodes are partitioned into remote components and each component is replaced with an internal `ModalUniversalExecutor_<hash>` proxy node.
|
||||
3. Referenced model assets are mirrored into storage.
|
||||
4. If custom-node syncing is enabled, the local `custom_nodes/` directory is zipped and mirrored if its content hash changed.
|
||||
|
||||
+38
-2
@@ -8,7 +8,7 @@ import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from typing import Any, Iterator
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -173,6 +173,42 @@ def _is_link(value: Any) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _looks_like_workflow_node(fragment: dict[str, Any]) -> bool:
|
||||
"""Return whether a JSON fragment resembles a saved ComfyUI workflow node."""
|
||||
return "id" in fragment and "properties" in fragment
|
||||
|
||||
|
||||
def _iter_workflow_nodes(
|
||||
workflow_fragment: Any,
|
||||
visited_object_ids: set[int] | None = None,
|
||||
) -> Iterator[dict[str, Any]]:
|
||||
"""Yield workflow node dictionaries from a nested saved workflow fragment."""
|
||||
if visited_object_ids is None:
|
||||
visited_object_ids = set()
|
||||
|
||||
if isinstance(workflow_fragment, dict):
|
||||
object_id = id(workflow_fragment)
|
||||
if object_id in visited_object_ids:
|
||||
return
|
||||
visited_object_ids.add(object_id)
|
||||
|
||||
if _looks_like_workflow_node(workflow_fragment):
|
||||
yield workflow_fragment
|
||||
|
||||
for value in workflow_fragment.values():
|
||||
yield from _iter_workflow_nodes(value, visited_object_ids)
|
||||
return
|
||||
|
||||
if isinstance(workflow_fragment, list):
|
||||
object_id = id(workflow_fragment)
|
||||
if object_id in visited_object_ids:
|
||||
return
|
||||
visited_object_ids.add(object_id)
|
||||
|
||||
for item in workflow_fragment:
|
||||
yield from _iter_workflow_nodes(item, visited_object_ids)
|
||||
|
||||
|
||||
def extract_remote_node_ids(
|
||||
workflow: dict[str, Any] | None,
|
||||
settings: ModalSyncSettings | None = None,
|
||||
@@ -183,7 +219,7 @@ def extract_remote_node_ids(
|
||||
|
||||
marker = (settings or get_settings()).marker_property
|
||||
remote_node_ids: set[str] = set()
|
||||
for node in workflow.get("nodes", []):
|
||||
for node in _iter_workflow_nodes(workflow):
|
||||
properties = node.get("properties") or {}
|
||||
if properties.get(marker):
|
||||
remote_node_ids.add(str(node.get("id")))
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
## 2026-04-17
|
||||
|
||||
- The workflow metadata already rides along in `extra_data.extra_pnginfo.workflow`, so the backend can discover `node.properties.is_modal_remote` without any extra frontend metadata channel.
|
||||
- That workflow metadata is not always flat. ComfyUI can embed saved node records inside nested subgraph fragments, so Modal marker discovery has to recurse through workflow JSON instead of only scanning the top-level `workflow["nodes"]` list.
|
||||
- A single wildcard-output proxy node is not enough for server-side graph rewriting. ComfyUI validates downstream link indices and output arity against the replacement class, so the backend has to generate proxy node classes that mirror the original node’s output signature before validation runs.
|
||||
- Keeping the sync engine content-addressable makes prompt rewrites cheap after the first upload. The route can block on sync safely as long as repeated requests collapse to a hash marker check.
|
||||
- Hashing and zipping the entire `custom_nodes/` tree inside the queue route is too expensive for `COMFY_MODAL_EXECUTION_MODE=local`. Local mode should skip custom-node bundle sync by default, with an explicit override when that path needs to be exercised.
|
||||
|
||||
@@ -153,6 +153,47 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy(
|
||||
assert summary.synced_assets[0].uploaded is True
|
||||
|
||||
|
||||
def test_extract_remote_node_ids_recurses_into_nested_subgraph_workflows(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
) -> None:
|
||||
"""Modal marker extraction should find nodes nested inside saved subgraph metadata."""
|
||||
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=Path("/tmp/storage"),
|
||||
remote_storage_root="/storage",
|
||||
custom_nodes_archive_name="custom_nodes_bundle.zip",
|
||||
comfyui_root=None,
|
||||
custom_nodes_dir=Path("/tmp/custom_nodes"),
|
||||
)
|
||||
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 100,
|
||||
"properties": {"is_modal_remote": False},
|
||||
"subgraph": {
|
||||
"nodes": [
|
||||
{"id": 11, "properties": {"is_modal_remote": True}},
|
||||
{"id": 12, "properties": {"is_modal_remote": False}},
|
||||
]
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
assert api_intercept_module.extract_remote_node_ids(workflow, settings) == {"11"}
|
||||
|
||||
|
||||
def test_rewrite_rejects_non_transportable_remote_inputs(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
@@ -226,6 +267,95 @@ def test_rewrite_rejects_non_transportable_remote_inputs(
|
||||
assert summary.remote_component_ids == ["1"]
|
||||
|
||||
|
||||
def test_rewrite_detects_remote_marker_inside_nested_subgraph_workflow(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
sync_engine_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Prompt rewrite should honor Modal markers found inside nested subgraph metadata."""
|
||||
custom_nodes_dir = tmp_path / "custom_nodes"
|
||||
custom_nodes_dir.mkdir()
|
||||
(custom_nodes_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8")
|
||||
|
||||
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=custom_nodes_dir,
|
||||
)
|
||||
sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
|
||||
fake_nodes_module = type(
|
||||
"FakeNodesModule",
|
||||
(),
|
||||
{
|
||||
"NODE_CLASS_MAPPINGS": {
|
||||
"ModelSource": _FakeRemoteModelNode,
|
||||
"RemoteConsumer": _FakeRemoteSamplerNode,
|
||||
"LocalConsumer": _FakeLocalSinkNode,
|
||||
},
|
||||
"NODE_DISPLAY_NAME_MAPPINGS": {},
|
||||
},
|
||||
)()
|
||||
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 99,
|
||||
"properties": {"is_modal_remote": False},
|
||||
"subgraph": {
|
||||
"nodes": [
|
||||
{"id": 1, "properties": {"is_modal_remote": False}},
|
||||
{"id": 2, "properties": {"is_modal_remote": True}},
|
||||
]
|
||||
},
|
||||
},
|
||||
{"id": 3, "properties": {"is_modal_remote": False}},
|
||||
]
|
||||
}
|
||||
prompt = {
|
||||
"1": {
|
||||
"class_type": "ModelSource",
|
||||
"inputs": {},
|
||||
"_meta": {"title": "Model Source"},
|
||||
},
|
||||
"2": {
|
||||
"class_type": "RemoteConsumer",
|
||||
"inputs": {"model": ["1", 0]},
|
||||
"_meta": {"title": "Remote Consumer"},
|
||||
},
|
||||
"3": {
|
||||
"class_type": "LocalConsumer",
|
||||
"inputs": {"latent": ["2", 0]},
|
||||
"_meta": {"title": "Local Consumer"},
|
||||
},
|
||||
}
|
||||
|
||||
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 list(rewritten_prompt) == ["1", "3"]
|
||||
assert rewritten_prompt["3"]["inputs"]["latent"] == ["1", 0]
|
||||
assert summary.remote_node_ids == ["1", "2"]
|
||||
assert summary.remote_component_ids == ["1"]
|
||||
|
||||
|
||||
def test_rewrite_auto_expands_upstream_non_transportable_dependencies(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
|
||||
Reference in New Issue
Block a user