feat: add modal remote expansion context menu
This commit is contained in:
Ken Simpson
@@ -138,6 +138,8 @@ For each connected region you want to offload:
|
||||
2. Enable it on every node that should be part of the same remote island.
|
||||
3. Confirm the node shows the blue remote-execution border.
|
||||
|
||||
If queue-time validation tells you a remote node still depends on local-only runtime objects, you no longer have to hunt those upstream nodes manually. Right-click the node and use `Modal: Include Required Upstream Nodes` to ask the backend which extra nodes must join that remote island. If multiple nodes are selected in the current graph, the same context menu expands the whole selection at once.
|
||||
|
||||
The toggle stores `properties.is_modal_remote = true` in workflow metadata. The visible graph is not rewritten in the editor. Rewrite happens only at queue time.
|
||||
|
||||
### 4. Queue the workflow
|
||||
|
||||
+291
-11
@@ -94,6 +94,33 @@ class RewriteSummary:
|
||||
custom_nodes_bundle: SyncedAsset | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RemoteExpansionReason:
|
||||
"""Describe why one upstream node had to join a remote component."""
|
||||
|
||||
node_id: str
|
||||
class_type: str
|
||||
required_by_node_id: str
|
||||
required_by_class_type: str
|
||||
output_index: int
|
||||
io_type: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class RemoteNodeAnalysis:
|
||||
"""Structured dry-run result for context-menu remote expansion."""
|
||||
|
||||
requested_node_ids: list[str] = field(default_factory=list)
|
||||
requested_workflow_node_paths: list[str] = field(default_factory=list)
|
||||
current_remote_node_ids: list[str] = field(default_factory=list)
|
||||
current_remote_workflow_node_paths: list[str] = field(default_factory=list)
|
||||
resolved_remote_node_ids: list[str] = field(default_factory=list)
|
||||
resolved_workflow_node_paths: list[str] = field(default_factory=list)
|
||||
added_node_ids: list[str] = field(default_factory=list)
|
||||
added_workflow_node_paths: list[str] = field(default_factory=list)
|
||||
reasons: list[RemoteExpansionReason] = field(default_factory=list)
|
||||
|
||||
|
||||
class ModalPromptValidationError(ValueError):
|
||||
"""Raised when a prompt cannot be executed with the current Modal transport."""
|
||||
|
||||
@@ -289,6 +316,118 @@ def _resolve_prompt_node_ids_for_workflow_node(
|
||||
return resolved_prompt_node_ids
|
||||
|
||||
|
||||
def _workflow_node_path(workflow_node_id: str, ancestor_node_ids: tuple[str, ...]) -> str:
|
||||
"""Return one workflow node's composed path, including subgraph ancestors."""
|
||||
if not ancestor_node_ids:
|
||||
return workflow_node_id
|
||||
return ":".join((*ancestor_node_ids, workflow_node_id))
|
||||
|
||||
|
||||
def _extract_marked_workflow_node_paths(
|
||||
workflow: dict[str, Any] | None,
|
||||
settings: ModalSyncSettings | None = None,
|
||||
) -> set[str]:
|
||||
"""Return composed workflow paths for nodes explicitly marked remote in metadata."""
|
||||
if workflow is None:
|
||||
return set()
|
||||
|
||||
marker = (settings or get_settings()).marker_property
|
||||
marked_workflow_node_paths: set[str] = set()
|
||||
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
||||
properties = node.get("properties") or {}
|
||||
if not properties.get(marker):
|
||||
continue
|
||||
marked_workflow_node_paths.add(
|
||||
_workflow_node_path(str(node.get("id")), ancestor_node_ids)
|
||||
)
|
||||
return marked_workflow_node_paths
|
||||
|
||||
|
||||
def _build_workflow_prompt_resolution_maps(
|
||||
workflow: dict[str, Any] | None,
|
||||
prompt_node_ids: set[str],
|
||||
) -> tuple[dict[str, set[str]], dict[str, set[str]]]:
|
||||
"""Return bidirectional mappings between workflow paths and queued prompt ids."""
|
||||
workflow_path_to_prompt_node_ids: dict[str, set[str]] = {}
|
||||
prompt_node_id_to_workflow_paths: dict[str, set[str]] = defaultdict(set)
|
||||
|
||||
if workflow is None:
|
||||
for prompt_node_id in prompt_node_ids:
|
||||
workflow_path_to_prompt_node_ids[prompt_node_id] = {prompt_node_id}
|
||||
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
||||
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
||||
|
||||
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
||||
workflow_node_id = str(node.get("id"))
|
||||
workflow_node_path = _workflow_node_path(workflow_node_id, ancestor_node_ids)
|
||||
resolved_prompt_node_ids = _resolve_prompt_node_ids_for_workflow_node(
|
||||
workflow_node_id,
|
||||
ancestor_node_ids,
|
||||
prompt_node_ids,
|
||||
)
|
||||
if workflow_node_id in prompt_node_ids:
|
||||
resolved_prompt_node_ids.add(workflow_node_id)
|
||||
if workflow_node_path in prompt_node_ids:
|
||||
resolved_prompt_node_ids.add(workflow_node_path)
|
||||
if not resolved_prompt_node_ids:
|
||||
continue
|
||||
|
||||
workflow_path_to_prompt_node_ids[workflow_node_path] = resolved_prompt_node_ids
|
||||
for prompt_node_id in resolved_prompt_node_ids:
|
||||
prompt_node_id_to_workflow_paths[prompt_node_id].add(workflow_node_path)
|
||||
|
||||
for prompt_node_id in prompt_node_ids:
|
||||
workflow_path_to_prompt_node_ids.setdefault(prompt_node_id, {prompt_node_id})
|
||||
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
||||
|
||||
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
||||
|
||||
|
||||
def _resolve_requested_prompt_node_ids(
|
||||
requested_workflow_node_paths: set[str],
|
||||
prompt_node_ids: set[str],
|
||||
workflow_path_to_prompt_node_ids: dict[str, set[str]],
|
||||
) -> set[str]:
|
||||
"""Resolve requested workflow-node paths to the prompt ids that queue-time rewrite sees."""
|
||||
requested_prompt_node_ids: set[str] = set()
|
||||
for requested_workflow_node_path in requested_workflow_node_paths:
|
||||
if requested_workflow_node_path in workflow_path_to_prompt_node_ids:
|
||||
requested_prompt_node_ids.update(
|
||||
workflow_path_to_prompt_node_ids[requested_workflow_node_path]
|
||||
)
|
||||
continue
|
||||
if requested_workflow_node_path in prompt_node_ids:
|
||||
requested_prompt_node_ids.add(requested_workflow_node_path)
|
||||
return requested_prompt_node_ids
|
||||
|
||||
|
||||
def _best_workflow_path_for_prompt_node(
|
||||
prompt_node_id: str,
|
||||
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
||||
) -> str:
|
||||
"""Choose the most specific workflow path for one queued prompt node id."""
|
||||
workflow_node_paths = prompt_node_id_to_workflow_paths.get(prompt_node_id, {prompt_node_id})
|
||||
return max(
|
||||
workflow_node_paths,
|
||||
key=lambda workflow_node_path: (
|
||||
workflow_node_path.count(":"),
|
||||
len(workflow_node_path),
|
||||
workflow_node_path,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _resolve_workflow_node_paths_for_prompt_nodes(
|
||||
prompt_node_ids: set[str],
|
||||
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
||||
) -> set[str]:
|
||||
"""Map queued prompt ids back to the workflow node paths the UI can mark remote."""
|
||||
return {
|
||||
_best_workflow_path_for_prompt_node(prompt_node_id, prompt_node_id_to_workflow_paths)
|
||||
for prompt_node_id in prompt_node_ids
|
||||
}
|
||||
|
||||
|
||||
def extract_remote_node_ids(
|
||||
workflow: dict[str, Any] | None,
|
||||
settings: ModalSyncSettings | None = None,
|
||||
@@ -404,10 +543,11 @@ def _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
prompt: dict[str, Any],
|
||||
remote_node_ids: set[str],
|
||||
nodes_module: Any,
|
||||
) -> set[str]:
|
||||
) -> tuple[set[str], list[RemoteExpansionReason]]:
|
||||
"""Grow the remote set upstream until non-transportable inputs stay inside the remote island."""
|
||||
expanded_remote_node_ids = set(remote_node_ids)
|
||||
added_node_ids: list[str] = []
|
||||
added_node_ids: set[str] = set()
|
||||
reasons: list[RemoteExpansionReason] = []
|
||||
|
||||
changed = True
|
||||
while changed:
|
||||
@@ -442,14 +582,23 @@ def _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
continue
|
||||
|
||||
expanded_remote_node_ids.add(upstream_node_id)
|
||||
added_node_ids.append(upstream_node_id)
|
||||
reason = RemoteExpansionReason(
|
||||
node_id=upstream_node_id,
|
||||
class_type=upstream_class_type,
|
||||
required_by_node_id=node_id,
|
||||
required_by_class_type=str(prompt_node["class_type"]),
|
||||
output_index=output_index,
|
||||
io_type=io_type,
|
||||
)
|
||||
reasons.append(reason)
|
||||
added_node_ids.add(upstream_node_id)
|
||||
changed = True
|
||||
logger.info(
|
||||
"Auto-expanded remote execution upstream: added node %s (%s) because node %s depends on non-transportable type '%s'.",
|
||||
upstream_node_id,
|
||||
upstream_class_type,
|
||||
node_id,
|
||||
io_type,
|
||||
reason.node_id,
|
||||
reason.class_type,
|
||||
reason.required_by_node_id,
|
||||
reason.io_type,
|
||||
)
|
||||
break
|
||||
if changed:
|
||||
@@ -460,9 +609,71 @@ def _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
"Expanded remote node set from %d to %d nodes by absorbing upstream non-transportable dependencies: %s",
|
||||
len(remote_node_ids),
|
||||
len(expanded_remote_node_ids),
|
||||
sorted(set(added_node_ids)),
|
||||
sorted(added_node_ids),
|
||||
)
|
||||
return expanded_remote_node_ids
|
||||
return expanded_remote_node_ids, reasons
|
||||
|
||||
|
||||
def analyze_remote_node_selection(
|
||||
prompt: dict[str, Any],
|
||||
workflow: dict[str, Any] | None,
|
||||
seed_workflow_node_paths: list[str],
|
||||
settings: ModalSyncSettings | None = None,
|
||||
nodes_module: Any | None = None,
|
||||
) -> RemoteNodeAnalysis:
|
||||
"""Return the nodes the UI should mark remote for one context-menu expansion request."""
|
||||
resolved_settings = settings or get_settings()
|
||||
resolved_nodes_module = nodes_module or _get_nodes_module()
|
||||
prompt_node_ids = {str(node_id) for node_id in prompt.keys()}
|
||||
requested_workflow_node_paths = {
|
||||
str(seed_workflow_node_path)
|
||||
for seed_workflow_node_path in seed_workflow_node_paths
|
||||
if str(seed_workflow_node_path)
|
||||
}
|
||||
current_remote_node_ids = extract_remote_node_ids(
|
||||
workflow,
|
||||
resolved_settings,
|
||||
prompt_node_ids,
|
||||
)
|
||||
current_remote_workflow_node_paths = _extract_marked_workflow_node_paths(
|
||||
workflow,
|
||||
resolved_settings,
|
||||
)
|
||||
workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths = (
|
||||
_build_workflow_prompt_resolution_maps(workflow, prompt_node_ids)
|
||||
)
|
||||
requested_node_ids = _resolve_requested_prompt_node_ids(
|
||||
requested_workflow_node_paths,
|
||||
prompt_node_ids,
|
||||
workflow_path_to_prompt_node_ids,
|
||||
)
|
||||
initial_remote_node_ids = current_remote_node_ids | requested_node_ids
|
||||
resolved_remote_node_ids, reasons = _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
prompt=prompt,
|
||||
remote_node_ids=initial_remote_node_ids,
|
||||
nodes_module=resolved_nodes_module,
|
||||
)
|
||||
resolved_workflow_node_paths = (
|
||||
_resolve_workflow_node_paths_for_prompt_nodes(
|
||||
resolved_remote_node_ids,
|
||||
prompt_node_id_to_workflow_paths,
|
||||
)
|
||||
| current_remote_workflow_node_paths
|
||||
)
|
||||
added_node_ids = resolved_remote_node_ids - current_remote_node_ids
|
||||
added_workflow_node_paths = resolved_workflow_node_paths - current_remote_workflow_node_paths
|
||||
|
||||
return RemoteNodeAnalysis(
|
||||
requested_node_ids=sorted(requested_node_ids),
|
||||
requested_workflow_node_paths=sorted(requested_workflow_node_paths),
|
||||
current_remote_node_ids=sorted(current_remote_node_ids),
|
||||
current_remote_workflow_node_paths=sorted(current_remote_workflow_node_paths),
|
||||
resolved_remote_node_ids=sorted(resolved_remote_node_ids),
|
||||
resolved_workflow_node_paths=sorted(resolved_workflow_node_paths),
|
||||
added_node_ids=sorted(added_node_ids),
|
||||
added_workflow_node_paths=sorted(added_workflow_node_paths),
|
||||
reasons=reasons,
|
||||
)
|
||||
|
||||
|
||||
def _build_component_plan(
|
||||
@@ -845,7 +1056,7 @@ def rewrite_prompt_for_modal(
|
||||
resolved_nodes_module = nodes_module or _get_nodes_module()
|
||||
resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings)
|
||||
rewritten_prompt = copy.deepcopy(prompt)
|
||||
expanded_remote_node_ids = _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
expanded_remote_node_ids, _ = _expand_remote_node_ids_for_non_transportable_inputs(
|
||||
prompt=rewritten_prompt,
|
||||
remote_node_ids=remote_node_ids,
|
||||
nodes_module=resolved_nodes_module,
|
||||
@@ -960,6 +1171,13 @@ async def _queue_prompt_json(prompt_server: Any, json_data: dict[str, Any]) -> w
|
||||
)
|
||||
|
||||
|
||||
def _analysis_route_path(route_path: str) -> str:
|
||||
"""Return the sibling HTTP route used for dry-run remote-node expansion."""
|
||||
if route_path.endswith("/queue_prompt"):
|
||||
return f"{route_path.removesuffix('/queue_prompt')}/analyze_remote_nodes"
|
||||
return f"{route_path.rstrip('/')}/analyze_remote_nodes"
|
||||
|
||||
|
||||
def setup_modal_queue_route(
|
||||
prompt_server: Any | None = None,
|
||||
sync_engine: ModalAssetSyncEngine | None = None,
|
||||
@@ -983,6 +1201,64 @@ def setup_modal_queue_route(
|
||||
return
|
||||
|
||||
resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings)
|
||||
analysis_route_path = _analysis_route_path(resolved_settings.route_path)
|
||||
|
||||
@prompt_server.routes.post(analysis_route_path)
|
||||
async def modal_analyze_remote_nodes(request: web.Request) -> web.Response:
|
||||
"""Analyze which workflow nodes should be marked remote for the current graph."""
|
||||
logger.info("Received Modal remote-node analysis request.")
|
||||
try:
|
||||
request_started_at = time.perf_counter()
|
||||
json_data = await request.json()
|
||||
prompt = json_data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Modal remote-node analysis requires a 'prompt' object.")
|
||||
|
||||
workflow = json_data.get("workflow")
|
||||
seed_node_ids = json_data.get("seed_node_ids") or []
|
||||
if not isinstance(seed_node_ids, list):
|
||||
raise ValueError("Modal remote-node analysis requires 'seed_node_ids' to be a list.")
|
||||
|
||||
analysis = analyze_remote_node_selection(
|
||||
prompt=prompt,
|
||||
workflow=workflow if isinstance(workflow, dict) else None,
|
||||
seed_workflow_node_paths=[str(seed_node_id) for seed_node_id in seed_node_ids],
|
||||
settings=resolved_settings,
|
||||
)
|
||||
logger.info(
|
||||
"Modal remote-node analysis finished in %.3fs with %d requested nodes and %d additions.",
|
||||
time.perf_counter() - request_started_at,
|
||||
len(analysis.requested_workflow_node_paths),
|
||||
len(analysis.added_workflow_node_paths),
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"requested_node_ids": analysis.requested_node_ids,
|
||||
"requested_workflow_node_paths": analysis.requested_workflow_node_paths,
|
||||
"current_remote_node_ids": analysis.current_remote_node_ids,
|
||||
"current_remote_workflow_node_paths": (
|
||||
analysis.current_remote_workflow_node_paths
|
||||
),
|
||||
"resolved_remote_node_ids": analysis.resolved_remote_node_ids,
|
||||
"resolved_workflow_node_paths": analysis.resolved_workflow_node_paths,
|
||||
"added_node_ids": analysis.added_node_ids,
|
||||
"added_workflow_node_paths": analysis.added_workflow_node_paths,
|
||||
"reasons": [
|
||||
{
|
||||
"node_id": reason.node_id,
|
||||
"class_type": reason.class_type,
|
||||
"required_by_node_id": reason.required_by_node_id,
|
||||
"required_by_class_type": reason.required_by_class_type,
|
||||
"output_index": reason.output_index,
|
||||
"io_type": reason.io_type,
|
||||
}
|
||||
for reason in analysis.reasons
|
||||
],
|
||||
}
|
||||
)
|
||||
except (TypeError, ValueError) as exc:
|
||||
logger.warning("Modal remote-node analysis request was invalid: %s", exc)
|
||||
return web.json_response({"error": str(exc), "node_errors": []}, status=400)
|
||||
|
||||
@prompt_server.routes.post(resolved_settings.route_path)
|
||||
async def modal_queue_prompt(request: web.Request) -> web.Response:
|
||||
@@ -1095,4 +1371,8 @@ def setup_modal_queue_route(
|
||||
return web.json_response({"error": str(exc), "node_errors": []}, status=500)
|
||||
|
||||
_ROUTE_REGISTERED = True
|
||||
logger.info("Registered Modal queue route at %s", resolved_settings.route_path)
|
||||
logger.info(
|
||||
"Registered Modal queue route at %s and analysis route at %s",
|
||||
resolved_settings.route_path,
|
||||
analysis_route_path,
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
- ComfyUI registry setup is mostly metadata, not code. The repo needs a real `[tool.comfy]` block in `pyproject.toml`, a stable immutable node-pack id, repository URLs, and an optional GitHub publish workflow; choosing a software license should stay explicit instead of being guessed during packaging work.
|
||||
- `comfy node pack` is not git-aware by default; it uses `.gitignore` as its practical packaging filter. For a clean registry artifact, this repo needs an allowlist-style `.gitignore` so untracked local files like logs, caches, virtualenvs, and scratch docs cannot leak into `node.zip`.
|
||||
- Remote per-node phase events are not enough to mirror sampler progress in the local UI. The Modal worker also needs a headless equivalent of ComfyUI's global progress hook so `ProgressBar` updates feed `progress_state`, and the local relay should forward those numeric updates as a separate high-frequency UI event rather than overloading coarse execution-phase messages.
|
||||
- Queue-time remote-boundary fixes should come from the backend's actual rewrite logic, not from parsing error text in the browser. A small dry-run analysis route can reuse the same prompt-id resolution and upstream auto-expansion rules, then return workflow-node paths for the frontend to mark directly from a context-menu action, including nested subgraph nodes.
|
||||
|
||||
## 2026-04-17
|
||||
|
||||
|
||||
@@ -582,6 +582,159 @@ def test_rewrite_auto_expands_upstream_non_transportable_dependencies(
|
||||
assert rewritten_prompt["4"]["inputs"]["latent"] == ["1", 0]
|
||||
|
||||
|
||||
def test_analyze_remote_node_selection_returns_nodes_to_mark_and_reasons(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
) -> None:
|
||||
"""Dry-run analysis should surface the clicked node plus required upstream nodes."""
|
||||
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"),
|
||||
)
|
||||
fake_nodes_module = type(
|
||||
"FakeNodesModule",
|
||||
(),
|
||||
{
|
||||
"NODE_CLASS_MAPPINGS": {
|
||||
"ModelSource": _FakeRemoteModelNode,
|
||||
"ConditioningSource": _FakeRemoteConditioningNode,
|
||||
"RemoteConsumer": _FakeRemoteSamplerNode,
|
||||
},
|
||||
"NODE_DISPLAY_NAME_MAPPINGS": {},
|
||||
},
|
||||
)()
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{"id": 1, "properties": {"is_modal_remote": False}},
|
||||
{"id": 2, "properties": {"is_modal_remote": False}},
|
||||
{"id": 3, "properties": {"is_modal_remote": False}},
|
||||
]
|
||||
}
|
||||
prompt = {
|
||||
"1": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}},
|
||||
"2": {
|
||||
"class_type": "ConditioningSource",
|
||||
"inputs": {},
|
||||
"_meta": {"title": "Conditioning"},
|
||||
},
|
||||
"3": {
|
||||
"class_type": "RemoteConsumer",
|
||||
"inputs": {"model": ["1", 0], "conditioning": ["2", 0]},
|
||||
"_meta": {"title": "Remote Consumer"},
|
||||
},
|
||||
}
|
||||
|
||||
analysis = api_intercept_module.analyze_remote_node_selection(
|
||||
prompt=prompt,
|
||||
workflow=workflow,
|
||||
seed_workflow_node_paths=["3"],
|
||||
settings=settings,
|
||||
nodes_module=fake_nodes_module,
|
||||
)
|
||||
|
||||
assert analysis.requested_node_ids == ["3"]
|
||||
assert analysis.requested_workflow_node_paths == ["3"]
|
||||
assert analysis.current_remote_node_ids == []
|
||||
assert analysis.current_remote_workflow_node_paths == []
|
||||
assert analysis.resolved_remote_node_ids == ["1", "2", "3"]
|
||||
assert analysis.resolved_workflow_node_paths == ["1", "2", "3"]
|
||||
assert analysis.added_node_ids == ["1", "2", "3"]
|
||||
assert analysis.added_workflow_node_paths == ["1", "2", "3"]
|
||||
assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [
|
||||
("1", "3"),
|
||||
("2", "3"),
|
||||
]
|
||||
|
||||
|
||||
def test_analyze_remote_node_selection_prefers_nested_workflow_paths(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
) -> None:
|
||||
"""Nested prompt ids should map back to the specific inner workflow node path."""
|
||||
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"),
|
||||
)
|
||||
fake_nodes_module = type(
|
||||
"FakeNodesModule",
|
||||
(),
|
||||
{
|
||||
"NODE_CLASS_MAPPINGS": {
|
||||
"ModelSource": _FakeRemoteModelNode,
|
||||
"RemoteConsumer": _FakeRemoteSamplerNode,
|
||||
},
|
||||
"NODE_DISPLAY_NAME_MAPPINGS": {},
|
||||
},
|
||||
)()
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 24,
|
||||
"properties": {"is_modal_remote": False},
|
||||
"subgraph": {
|
||||
"nodes": [
|
||||
{"id": 23, "properties": {"is_modal_remote": True}},
|
||||
]
|
||||
},
|
||||
},
|
||||
{"id": 30, "properties": {"is_modal_remote": False}},
|
||||
]
|
||||
}
|
||||
prompt = {
|
||||
"30": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}},
|
||||
"24:23": {
|
||||
"class_type": "RemoteConsumer",
|
||||
"inputs": {"model": ["30", 0]},
|
||||
"_meta": {"title": "Nested Consumer"},
|
||||
},
|
||||
}
|
||||
|
||||
analysis = api_intercept_module.analyze_remote_node_selection(
|
||||
prompt=prompt,
|
||||
workflow=workflow,
|
||||
seed_workflow_node_paths=["24:23"],
|
||||
settings=settings,
|
||||
nodes_module=fake_nodes_module,
|
||||
)
|
||||
|
||||
assert analysis.requested_node_ids == ["24:23"]
|
||||
assert analysis.current_remote_node_ids == ["24:23"]
|
||||
assert analysis.current_remote_workflow_node_paths == ["24:23"]
|
||||
assert analysis.resolved_remote_node_ids == ["24:23", "30"]
|
||||
assert analysis.resolved_workflow_node_paths == ["24:23", "30"]
|
||||
assert analysis.added_node_ids == ["30"]
|
||||
assert analysis.added_workflow_node_paths == ["30"]
|
||||
assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [
|
||||
("30", "24:23"),
|
||||
]
|
||||
|
||||
|
||||
def test_rewrite_rejects_non_transportable_remote_outputs(
|
||||
api_intercept_module: Any,
|
||||
settings_module: Any,
|
||||
|
||||
@@ -87,6 +87,32 @@ def test_streamed_modal_node_progress_updates_active_overlay() -> None:
|
||||
assert 'ctx.fillText(`${Math.round(progressRatio * 100)}%`' in source
|
||||
|
||||
|
||||
def test_modal_context_menu_can_expand_required_upstream_nodes() -> None:
|
||||
"""Right-clicking a node should offer a dry-run expansion action backed by the backend."""
|
||||
source = _modal_toggle_source()
|
||||
|
||||
assert 'const MODAL_ANALYZE_ROUTE = MODAL_ROUTE.replace(/\\/queue_prompt$/, "/analyze_remote_nodes");' in source
|
||||
assert "function workflowNodePath(node)" in source
|
||||
assert "function findNodeByWorkflowPath(workflowPath)" in source
|
||||
assert "function selectedWorkflowNodePaths(node)" in source
|
||||
assert "function analyzeAndMarkRequiredRemoteNodes(node)" in source
|
||||
assert 'api.fetchApi(MODAL_ANALYZE_ROUTE, {' in source
|
||||
assert '"Modal: Include Required Upstream Nodes"' in source
|
||||
assert '"Modal: Include Required Upstream Nodes for Selection"' in source
|
||||
|
||||
|
||||
def test_modal_context_menu_marks_nodes_across_subgraphs() -> None:
|
||||
"""The UI action should be able to resolve and mark nested workflow-node paths."""
|
||||
source = _modal_toggle_source()
|
||||
|
||||
assert "function rootGraph()" in source
|
||||
assert "function findSomethingInAllSubgraphs(matcher)" in source
|
||||
assert "function findContainingSubgraphNode(subgraphId)" in source
|
||||
assert "candidate.subgraph?.id === subgraphId" in source
|
||||
assert "function markWorkflowNodePathsRemote(workflowNodePaths)" in source
|
||||
assert "setRemoteFlag(node, true);" in source
|
||||
|
||||
|
||||
def test_prompt_cleanup_prunes_orphaned_global_status_entries() -> None:
|
||||
"""Workflow cleanup should remove stale global badge states once prompt state is gone."""
|
||||
source = _modal_toggle_source()
|
||||
|
||||
@@ -3,6 +3,7 @@ import { PromptExecutionError, api } from "../../scripts/api.js";
|
||||
|
||||
const REMOTE_PROPERTY = "is_modal_remote";
|
||||
const MODAL_ROUTE = "/modal/queue_prompt";
|
||||
const MODAL_ANALYZE_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/analyze_remote_nodes");
|
||||
const INTERNAL_NODE_PREFIX = "ModalUniversalExecutor";
|
||||
|
||||
const IDLE_BORDER_COLOR = "#1d9bf0";
|
||||
@@ -643,6 +644,229 @@ function extractRemoteNodeIds(workflow) {
|
||||
return remoteNodeIds;
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the root workflow graph, including subgraphs when available.
|
||||
* @returns {LGraph | null}
|
||||
*/
|
||||
function rootGraph() {
|
||||
return app.rootGraph ?? app.graph?.rootGraph ?? app.graph ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Look up a graph-local node id without assuming whether ids are numeric or strings.
|
||||
* @param {LGraph | null | undefined} graph
|
||||
* @param {string} id
|
||||
* @returns {LGraphNode | null}
|
||||
*/
|
||||
function getGraphNodeById(graph, id) {
|
||||
if (!graph || id == null) {
|
||||
return null;
|
||||
}
|
||||
const directMatch = graph.getNodeById?.(id) ?? null;
|
||||
if (directMatch) {
|
||||
return directMatch;
|
||||
}
|
||||
const numericId = Number(id);
|
||||
if (Number.isFinite(numericId)) {
|
||||
return graph.getNodeById?.(numericId) ?? null;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Search every live subgraph for one matching value.
|
||||
* @param {(graph: LGraph) => any} matcher
|
||||
* @returns {any}
|
||||
*/
|
||||
function findSomethingInAllSubgraphs(matcher) {
|
||||
const graph = rootGraph();
|
||||
if (!graph) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const subgraphs = [graph];
|
||||
if (graph.subgraphs?.values) {
|
||||
subgraphs.push(...graph.subgraphs.values());
|
||||
}
|
||||
for (const subgraph of subgraphs) {
|
||||
const match = matcher(subgraph);
|
||||
if (match) {
|
||||
return match;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the workflow node that owns one nested subgraph graph id.
|
||||
* @param {number | string | undefined} subgraphId
|
||||
* @returns {LGraphNode | null}
|
||||
*/
|
||||
function findContainingSubgraphNode(subgraphId) {
|
||||
if (subgraphId == null) {
|
||||
return null;
|
||||
}
|
||||
return (
|
||||
findSomethingInAllSubgraphs((graph) =>
|
||||
(graph?.nodes ?? []).find(
|
||||
(candidate) =>
|
||||
typeof candidate?.isSubgraphNode === "function" &&
|
||||
candidate.isSubgraphNode() &&
|
||||
candidate.subgraph?.id === subgraphId,
|
||||
),
|
||||
) ?? null
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return one node's composed workflow path, including any subgraph ancestors.
|
||||
* @param {LGraphNode} node
|
||||
* @returns {string}
|
||||
*/
|
||||
function workflowNodePath(node) {
|
||||
const pathSegments = [String(node?.id ?? "")];
|
||||
let currentGraph = node?.graph ?? null;
|
||||
const currentRootGraph = rootGraph();
|
||||
|
||||
while (currentGraph && currentRootGraph && currentGraph !== currentRootGraph) {
|
||||
const parentNode = findContainingSubgraphNode(currentGraph.id);
|
||||
if (!parentNode) {
|
||||
break;
|
||||
}
|
||||
pathSegments.unshift(String(parentNode.id));
|
||||
currentGraph = parentNode.graph ?? null;
|
||||
}
|
||||
|
||||
return pathSegments.filter(Boolean).join(":");
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve a composed workflow path like `24:23` to a live node instance.
|
||||
* @param {string} workflowPath
|
||||
* @returns {LGraphNode | null}
|
||||
*/
|
||||
function findNodeByWorkflowPath(workflowPath) {
|
||||
const pathSegments = String(workflowPath)
|
||||
.split(":")
|
||||
.filter(Boolean);
|
||||
if (pathSegments.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
let currentGraph = rootGraph();
|
||||
let currentNode = null;
|
||||
for (const pathSegment of pathSegments) {
|
||||
currentNode = getGraphNodeById(currentGraph, pathSegment);
|
||||
if (!currentNode) {
|
||||
return null;
|
||||
}
|
||||
currentGraph = currentNode.subgraph ?? null;
|
||||
}
|
||||
return currentNode;
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the workflow-node paths that a context-menu action should expand from.
|
||||
* @param {LGraphNode} node
|
||||
* @returns {string[]}
|
||||
*/
|
||||
function selectedWorkflowNodePaths(node) {
|
||||
const selectedNodes = Object.values(app.canvas?.selected_nodes ?? {}).filter(
|
||||
(candidate) => candidate?.graph === node?.graph && isEligibleNode(candidate),
|
||||
);
|
||||
if (selectedNodes.some((candidate) => candidate === node) && selectedNodes.length > 1) {
|
||||
return selectedNodes.map((candidate) => workflowNodePath(candidate));
|
||||
}
|
||||
return [workflowNodePath(node)];
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the graph snapshot shape ComfyUI already uses for queue submission.
|
||||
* @returns {Promise<{ output: object, workflow: object }>}
|
||||
*/
|
||||
async function serializeCurrentGraphForModal() {
|
||||
if (typeof app.graphToPrompt !== "function") {
|
||||
throw new Error("ComfyUI graph serialization is unavailable.");
|
||||
}
|
||||
const prompt = await app.graphToPrompt();
|
||||
if (!prompt?.output || !prompt?.workflow) {
|
||||
throw new Error("ComfyUI did not return prompt and workflow data.");
|
||||
}
|
||||
return {
|
||||
output: prompt.output,
|
||||
workflow: prompt.workflow,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Show a short frontend notification without taking over the whole UI.
|
||||
* @param {string} value
|
||||
*/
|
||||
function notifyModal(value) {
|
||||
dispatchSyntheticApiEvent("notification", {
|
||||
id: `modal-analysis-${Date.now()}`,
|
||||
value,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply the remote marker to the workflow nodes named by composed workflow paths.
|
||||
* @param {string[]} workflowNodePaths
|
||||
* @returns {number}
|
||||
*/
|
||||
function markWorkflowNodePathsRemote(workflowNodePaths) {
|
||||
let appliedCount = 0;
|
||||
for (const workflowPath of workflowNodePaths) {
|
||||
const node = findNodeByWorkflowPath(workflowPath);
|
||||
if (!node) {
|
||||
console.warn("Unable to find Modal workflow node path in the live graph.", workflowPath);
|
||||
continue;
|
||||
}
|
||||
setRemoteFlag(node, true);
|
||||
appliedCount += 1;
|
||||
}
|
||||
return appliedCount;
|
||||
}
|
||||
|
||||
/**
|
||||
* Request required upstream nodes from the backend and mark them remote in the UI.
|
||||
* @param {LGraphNode} node
|
||||
*/
|
||||
async function analyzeAndMarkRequiredRemoteNodes(node) {
|
||||
const seedNodeIds = selectedWorkflowNodePaths(node);
|
||||
const graphSnapshot = await serializeCurrentGraphForModal();
|
||||
const response = await api.fetchApi(MODAL_ANALYZE_ROUTE, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
prompt: graphSnapshot.output,
|
||||
workflow: graphSnapshot.workflow,
|
||||
seed_node_ids: seedNodeIds,
|
||||
}),
|
||||
});
|
||||
if (response.status !== 200) {
|
||||
throw new PromptExecutionError(await response.json());
|
||||
}
|
||||
|
||||
const result = await response.json();
|
||||
const resolvedWorkflowNodePaths = result.resolved_workflow_node_paths ?? [];
|
||||
const addedWorkflowNodePaths = result.added_workflow_node_paths ?? [];
|
||||
const appliedCount = markWorkflowNodePathsRemote(resolvedWorkflowNodePaths);
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
|
||||
if (addedWorkflowNodePaths.length > 0) {
|
||||
notifyModal(`Marked ${addedWorkflowNodePaths.length} node${addedWorkflowNodePaths.length === 1 ? "" : "s"} for Modal.`);
|
||||
return;
|
||||
}
|
||||
notifyModal(
|
||||
appliedCount > 0
|
||||
? "No extra upstream Modal nodes were required."
|
||||
: "Modal analysis finished, but no matching live nodes were found to update.",
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Update the node state and redraw the canvas.
|
||||
* @param {LGraphNode} node
|
||||
@@ -799,6 +1023,34 @@ function decorateNode(node) {
|
||||
originalDrawForeground?.apply(this, arguments);
|
||||
drawRemoteNodeDecoration(this, ctx);
|
||||
};
|
||||
|
||||
const originalGetExtraMenuOptions = node.getExtraMenuOptions;
|
||||
node.getExtraMenuOptions = function getExtraMenuOptionsWithModalAnalysis(_, options) {
|
||||
const menuOptions = originalGetExtraMenuOptions?.apply(this, arguments) ?? options ?? [];
|
||||
const targetOptions = options ?? menuOptions;
|
||||
if (!Array.isArray(targetOptions)) {
|
||||
return menuOptions;
|
||||
}
|
||||
|
||||
const selectedNodePaths = selectedWorkflowNodePaths(this);
|
||||
const menuItemLabel =
|
||||
selectedNodePaths.length > 1
|
||||
? "Modal: Include Required Upstream Nodes for Selection"
|
||||
: "Modal: Include Required Upstream Nodes";
|
||||
if (!targetOptions.some((option) => option?.content === menuItemLabel)) {
|
||||
targetOptions.push(null, {
|
||||
content: menuItemLabel,
|
||||
callback: () => {
|
||||
void analyzeAndMarkRequiredRemoteNodes(this).catch((error) => {
|
||||
console.error("Modal remote-node analysis failed.", error);
|
||||
notifyModal(`Modal remote-node analysis failed: ${String(error?.message ?? error)}`);
|
||||
});
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
return menuOptions;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user