diff --git a/README.md b/README.md index 6e54d8a..7603f8a 100644 --- a/README.md +++ b/README.md @@ -127,9 +127,10 @@ When a prompt is queued: 4. Components also expand across non-transportable inputs such as `MODEL`, `CLIP`, `VAE`, or `CONDITIONING`. 5. Each component is replaced with one or more generated `ModalUniversalExecutor_` proxy nodes. 6. Referenced model assets and, when enabled, `custom_nodes/` packages are mirrored into storage. -7. The rewritten prompt is submitted to ComfyUI's normal execution queue. -8. Local nodes execute normally until a proxy node is reached. -9. The proxy dispatches local or Modal execution. Values needed by local nodes are materialized as normal ComfyUI outputs, while values consumed only by later remote components travel through the local graph as small Modal-backed references. +7. An internal `ModalArtifactFinalizer` output sink is connected to every generated proxy so ComfyUI executes remote components even when their terminal nodes only save files and expose no normal ComfyUI output. +8. The rewritten prompt is submitted to ComfyUI's normal execution queue. +9. Local nodes execute normally until a proxy node is reached. +10. The proxy dispatches local or Modal execution. Values needed by local nodes are materialized as normal ComfyUI outputs, while values consumed only by later remote components travel through the local graph as small Modal-backed references. Boundary-crossing values must be transportable. Supported evaluated values include: @@ -177,6 +178,8 @@ Each remote payload also carries a stable invocation id. The worker records its After every remote payload, including the final component that clears its remote session, Modal-Sync compares the remote ComfyUI `output/` tree with its pre-execution snapshot. New or replaced regular files are bundled into the durable result and downloaded into the corresponding local `output/` subdirectory. Each downloaded filename is prefixed `remote---`, where `app_id` is the unique Modal app suffix and `epoch` is the trailing nine digits of the remote completion time. Artifact paths and SHA-256 digests are validated, symlinks and escaping paths are rejected, identical retry downloads are reused, and a differing local collision receives a numeric suffix instead of being overwritten. This lets a remote video-save node return its compressed file without transporting the decoded frame tensor back through ComfyUI. +The proxy emits its internal completion token only after the remote result and its artifacts have been materialized locally. The finalizer consumes those tokens without producing a user-visible value, making artifact-only remote branches part of ComfyUI's required output path while preserving normal downstream graph behavior. + Remote boundary hydration does not alter the prompt metadata exposed to custom nodes. The PromptExecutor uses hydrated tensors and runtime objects for dependency execution, while hidden `PROMPT` inputs receive a separately preserved JSON-safe graph. Metadata-writing output nodes such as VideoHelperSuite can therefore serialize the prompt into an image or video without encountering tensor values introduced by the remote boundary. Remote subgraph runs can persist transport-safe node outputs into a shared Modal `Dict` using ComfyUI input-signature semantics. The cache skips non-serializable outputs and entries above the configured size cap. diff --git a/__init__.py b/__init__.py index 44e0ea5..41b89e0 100644 --- a/__init__.py +++ b/__init__.py @@ -45,11 +45,16 @@ try: from comfy_api.latest import ComfyExtension, io from .api_intercept import setup_modal_queue_route - from .modal_executor_node import ModalMapInput, ModalUniversalExecutor + from .modal_executor_node import ( + ModalArtifactFinalizer, + ModalMapInput, + ModalUniversalExecutor, + ) except ModuleNotFoundError: # pragma: no cover - used during local non-Comfy imports. ComfyExtension = object # type: ignore[assignment] io = None # type: ignore[assignment] ModalMapInput = None # type: ignore[assignment] + ModalArtifactFinalizer = None # type: ignore[assignment] ModalUniversalExecutor = None # type: ignore[assignment] class ComfyModalSyncExtension: # type: ignore[no-redef] @@ -69,7 +74,7 @@ else: async def get_node_list(self) -> list[type[io.ComfyNode]]: """Return the node list after registering API routes.""" setup_modal_queue_route() - return [ModalUniversalExecutor, ModalMapInput] + return [ModalUniversalExecutor, ModalMapInput, ModalArtifactFinalizer] async def comfy_entrypoint() -> ComfyExtension: """Create the ComfyUI v3 extension entrypoint.""" @@ -79,6 +84,7 @@ else: __all__ = [ "ComfyModalSyncExtension", + "ModalArtifactFinalizer", "ModalMapInput", "ModalUniversalExecutor", "WEB_DIRECTORY", diff --git a/api_intercept.py b/api_intercept.py index 43a4f1c..4684f9d 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -17,7 +17,11 @@ from typing import Any, Iterator, Mapping from aiohttp import web from .modal_executor_node import ( + MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, + MODAL_ARTIFACT_FINALIZER_NODE_ID, + MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, MODAL_MAP_INPUT_NODE_ID, + ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, @@ -178,6 +182,7 @@ class RewriteSummary: rewritten_node_id_map: dict[str, str] = field(default_factory=dict) synced_assets: list[SyncedAsset] = field(default_factory=list) custom_nodes_bundle: SyncedAsset | None = None + artifact_finalizer_node_id: str | None = None @dataclass(frozen=True) @@ -2003,10 +2008,16 @@ def _build_component_plan( boundary_inputs_by_source[source] = spec spec.targets.append(InputTarget(node_id=node_id, input_name=str(input_name))) + has_downstream_consumer = False for output_index, io_type in enumerate(output_types): source = LinkedOutputRef(node_id=node_id, output_index=output_index) + output_consumers = consumers.get(source, []) + if output_consumers: + has_downstream_consumer = True local_consumers = [ - consumer for consumer in consumers.get(source, []) if consumer.node_id not in component_node_id_set + consumer + for consumer in output_consumers + if consumer.node_id not in component_node_id_set ] if not local_consumers: continue @@ -2025,6 +2036,9 @@ def _build_component_plan( ), ) + if not has_downstream_consumer: + output_execution_targets.add(node_id) + mapped_boundary_spec: BoundaryInputSpec | None = None mapped_boundary_input_io_type: str | None = None for boundary_input in boundary_inputs_by_source.values(): @@ -3398,6 +3412,7 @@ def _rewrite_component_into_proxy( output_is_list=tuple(bool(output.get("is_list", False)) for output in boundary_outputs), nodes_module=nodes_module, is_output_node=is_output_node, + include_completion_output=True, ) proxy_inputs["original_node_data"] = register_cache_friendly_proxy_payload( prompt_node_id, @@ -3599,6 +3614,7 @@ def _rewrite_component_into_proxy( output_is_list=tuple(spec.is_list for spec in component.boundary_outputs), nodes_module=nodes_module, is_output_node=component.contains_output_node, + include_completion_output=True, ) representative_node_id = component.representative_node_id proxy_inputs = proxy_inputs_for_boundary_inputs(list(payload.get("boundary_inputs", []))) @@ -3642,6 +3658,79 @@ def _rewrite_component_into_proxy( return [representative_node_id] +def _modal_component_completion_output_index( + *, + proxy_node_id: str, + rewritten_prompt: dict[str, Any], + nodes_module: Any, +) -> int: + """Return the synthetic completion-token output index for one Modal proxy.""" + prompt_node = rewritten_prompt.get(proxy_node_id) + if prompt_node is None: + raise ModalPromptValidationError( + f"Modal proxy node {proxy_node_id!r} is missing before artifact finalization." + ) + class_type = str(prompt_node.get("class_type")) + node_class = nodes_module.NODE_CLASS_MAPPINGS.get(class_type) + if node_class is None: + raise ModalPromptValidationError( + f"Modal proxy class {class_type!r} is not registered before artifact finalization." + ) + _output_types, output_names, _output_is_list = _normalize_output_metadata(node_class) + try: + return output_names.index(MODAL_COMPONENT_COMPLETION_OUTPUT_NAME) + except ValueError as exc: + raise ModalPromptValidationError( + f"Modal proxy {proxy_node_id!r} does not expose its completion token." + ) from exc + + +def _attach_modal_artifact_finalizer( + *, + rewritten_prompt: dict[str, Any], + remote_component_ids: list[str], + nodes_module: Any, +) -> str: + """Attach one internal output sink to every rewritten Modal component.""" + if not remote_component_ids: + raise ModalPromptValidationError( + "Modal artifact finalization requires at least one remote component." + ) + if len(remote_component_ids) > MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS: + raise ModalPromptValidationError( + "Modal artifact finalization supports at most " + f"{MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS} remote components per prompt." + ) + + ensure_modal_artifact_finalizer_registered(nodes_module) + finalizer_node_id = f"__{MODAL_ARTIFACT_FINALIZER_NODE_ID}__" + while finalizer_node_id in rewritten_prompt: + finalizer_node_id = f"{finalizer_node_id}_proxy" + + finalizer_inputs = { + f"component_{component_index}": [ + proxy_node_id, + _modal_component_completion_output_index( + proxy_node_id=proxy_node_id, + rewritten_prompt=rewritten_prompt, + nodes_module=nodes_module, + ), + ] + for component_index, proxy_node_id in enumerate(remote_component_ids) + } + rewritten_prompt[finalizer_node_id] = { + "class_type": MODAL_ARTIFACT_FINALIZER_NODE_ID, + "inputs": finalizer_inputs, + "_meta": {"title": "Modal Artifact Finalizer"}, + } + logger.info( + "Attached Modal artifact finalizer %s to remote components %s.", + finalizer_node_id, + remote_component_ids, + ) + return finalizer_node_id + + def rewrite_prompt_for_modal( prompt: dict[str, Any], workflow: dict[str, Any] | None, @@ -3820,6 +3909,12 @@ def rewrite_prompt_for_modal( if component.mapped_boundary_input_name: mapped_proxy_component_ids.add(proxy_node_ids[0]) + summary.artifact_finalizer_node_id = _attach_modal_artifact_finalizer( + rewritten_prompt=rewritten_prompt, + remote_component_ids=summary.remote_component_ids, + nodes_module=resolved_nodes_module, + ) + proxy_component_groups = { component_id: {component_id} for component_id in summary.remote_component_ids diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index b56c131..f99e6f4 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -4,6 +4,7 @@ - Iteration-rate labels can be derived entirely from successive streamed `value`/`max` samples in the frontend, but refocus replay must use the backend event timestamps rather than replay arrival time or buffered events will produce artificially high `it/s` values. Track and smooth rates independently for ordinary node progress and every mapped worker lane. - A fast prompt can emit `execution_success` and clear its frontend state before the queue route's successful HTTP response is processed. Queue-response UI reconciliation must treat a `null` prompt state as an already-terminal success and avoid recreating ready/waiting visuals or converting completion into a frontend `TypeError`. +- Artifact-only remote workflows need two execution anchors: the remote component planner must keep a terminal side-effect node in `execute_node_ids` even when it has no outputs, and the rewritten local prompt must expose a real ComfyUI output path. A synthetic completion token emitted after artifact materialization and consumed by an internal V3 output sink satisfies both requirements without adding a user-visible workflow output. ## 2026-08-07 diff --git a/modal_executor_node.py b/modal_executor_node.py index d63c3c8..7316b98 100644 --- a/modal_executor_node.py +++ b/modal_executor_node.py @@ -19,6 +19,9 @@ from .serialization import deserialize_node_outputs, serialize_node_inputs, spli logger = logging.getLogger(__name__) MODAL_MAP_INPUT_NODE_ID = "ModalMapInput" +MODAL_ARTIFACT_FINALIZER_NODE_ID = "ModalArtifactFinalizer" +MODAL_COMPONENT_COMPLETION_OUTPUT_NAME = "modal_component_complete" +MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS = 100 _PROXY_CACHE_CONTEXT_ID_KEY = "__comfy_modal_proxy_cache_context_id__" _VOLATILE_PROXY_CACHE_KEYS = frozenset( { @@ -497,6 +500,7 @@ def _build_proxy_node_class( output_is_list: tuple[bool, ...], *, is_output_node: bool, + include_completion_output: bool = False, ) -> type[io.ComfyNode]: """Create a v3 proxy node that mirrors an original node output signature.""" @@ -512,6 +516,12 @@ def _build_proxy_node_class( _output_spec(io_type, name, is_list) for io_type, name, is_list in zip(output_types, output_names, output_is_list, strict=False) ] + if include_completion_output: + outputs.append( + io.Boolean.Output( + display_name=MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, + ) + ) return io.Schema( node_id=node_id, display_name=proxy_display_name, @@ -553,6 +563,8 @@ def _build_proxy_node_class( payload.get("payload_kind"), len(outputs), ) + if include_completion_output: + return io.NodeOutput(*outputs, True) return io.NodeOutput(*outputs) _DynamicModalExecutor.__name__ = f"DynamicModalExecutor_{node_id}" @@ -596,6 +608,7 @@ def ensure_modal_component_proxy_node_registered( nodes_module: Any, *, is_output_node: bool, + include_completion_output: bool = False, ) -> str: """Register and return a proxy node id for a remote component signature.""" normalized_output_types = tuple(output_types) @@ -603,7 +616,8 @@ def ensure_modal_component_proxy_node_registered( normalized_output_is_list = tuple(output_is_list) proxy_node_id = _proxy_node_id( "ModalRemoteComponent", - normalized_output_types + (str(is_output_node),), + normalized_output_types + + (str(is_output_node), str(include_completion_output)), ) if proxy_node_id in _PROXY_NODE_CACHE: @@ -619,6 +633,7 @@ def ensure_modal_component_proxy_node_registered( output_names=normalized_output_names, output_is_list=normalized_output_is_list, is_output_node=is_output_node, + include_completion_output=include_completion_output, ) nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = proxy_class nodes_module.NODE_DISPLAY_NAME_MAPPINGS[proxy_node_id] = "Modal Remote Component" @@ -627,6 +642,16 @@ def ensure_modal_component_proxy_node_registered( return proxy_node_id +def ensure_modal_artifact_finalizer_registered(nodes_module: Any) -> None: + """Register the internal artifact-finalization sink in a ComfyUI node mapping.""" + nodes_module.NODE_CLASS_MAPPINGS[MODAL_ARTIFACT_FINALIZER_NODE_ID] = ( + ModalArtifactFinalizer + ) + nodes_module.NODE_DISPLAY_NAME_MAPPINGS[MODAL_ARTIFACT_FINALIZER_NODE_ID] = ( + "Modal Artifact Finalizer" + ) + + class ModalUniversalExecutor(io.ComfyNode): """Base debug node for Modal execution routing.""" @@ -686,3 +711,55 @@ class ModalMapInput(io.ComfyNode): value, ) return io.NodeOutput(value) + + +class ModalArtifactFinalizer(io.ComfyNode): + """Internal output sink that anchors remote execution and artifact materialization.""" + + @classmethod + def define_schema(cls) -> io.Schema: + """Accept one completion token from every rewritten terminal Modal component.""" + completion_template = io.Autogrow.TemplatePrefix( + input=io.Boolean.Input("completion", force_input=True), + prefix="component_", + min=1, + max=MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, + ) + return io.Schema( + node_id=MODAL_ARTIFACT_FINALIZER_NODE_ID, + display_name="Modal Artifact Finalizer", + category="Modal", + description=( + "Internal output sink inserted into rewritten prompts so Modal components " + "run and materialize their remote output artifacts locally." + ), + inputs=[ + io.Autogrow.Input( + "components", + template=completion_template, + ) + ], + outputs=[], + is_output_node=True, + is_dev_only=True, + is_experimental=True, + ) + + @classmethod + def execute(cls, components: io.Autogrow.Type) -> io.NodeOutput: + """Confirm every upstream Modal component completed before the prompt finishes.""" + incomplete_component_names = [ + str(component_name) + for component_name, completed in components.items() + if completed is not True + ] + if incomplete_component_names: + raise RuntimeError( + "Modal artifact finalization received incomplete component tokens: " + f"{incomplete_component_names}." + ) + logger.info( + "Modal artifact finalization completed after %d remote component(s).", + len(components), + ) + return io.NodeOutput() diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index eb68449..0b49919 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -141,6 +141,23 @@ class _FakePreviewImageNode: OUTPUT_NODE = True +class _FakeRemoteArtifactWriterNode: + """Fake terminal remote node that saves artifacts without ComfyUI outputs.""" + + RETURN_TYPES = () + RETURN_NAMES = () + OUTPUT_IS_LIST = () + OUTPUT_NODE = False + + +def _artifact_finalizer_node_id(summary: Any) -> str: + """Return the finalizer id after asserting that prompt rewrite attached it.""" + finalizer_node_id = summary.artifact_finalizer_node_id + assert isinstance(finalizer_node_id, str) + assert finalizer_node_id + return finalizer_node_id + + def test_rewrite_remote_mode_rejects_local_sync_backend( api_intercept_module: Any, settings_module: Any, @@ -907,7 +924,7 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy( extra_data={"extra_pnginfo": {"workflow": workflow}}, ) - assert set(rewritten_prompt) == {"1", "3"} + assert set(rewritten_prompt) == {"1", "3", _artifact_finalizer_node_id(summary)} rewritten_node = rewritten_prompt["1"] payload = rewritten_node["inputs"]["original_node_data"] assert rewritten_node["class_type"].startswith("ModalUniversalExecutor_") @@ -940,6 +957,91 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy( assert summary.synced_assets[0].uploaded is True +def test_rewrite_anchors_terminal_artifact_only_remote_node( + api_intercept_module: Any, + modal_executor_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """A terminal remote side-effect node should remain executable through the finalizer.""" + 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": { + "RemoteArtifactWriter": _FakeRemoteArtifactWriterNode, + }, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + workflow = {"nodes": [{"id": 1, "properties": {"is_modal_remote": True}}]} + prompt = { + "1": { + "class_type": "RemoteArtifactWriter", + "inputs": {}, + "_meta": {"title": "Remote Artifact Writer"}, + } + } + + 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, + ) + + finalizer_node_id = _artifact_finalizer_node_id(summary) + assert set(rewritten_prompt) == {"1", finalizer_node_id} + payload = rewritten_prompt["1"]["inputs"]["original_node_data"] + assert payload["execute_node_ids"] == ["1"] + assert payload["boundary_outputs"] == [] + + proxy_class_type = rewritten_prompt["1"]["class_type"] + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_class_type] + proxy_schema = proxy_class.GET_SCHEMA() + assert [output.io_type for output in proxy_schema.outputs] == ["BOOLEAN"] + assert [output.display_name for output in proxy_schema.outputs] == [ + modal_executor_module.MODAL_COMPONENT_COMPLETION_OUTPUT_NAME + ] + + assert rewritten_prompt[finalizer_node_id] == { + "class_type": modal_executor_module.MODAL_ARTIFACT_FINALIZER_NODE_ID, + "inputs": {"component_0": ["1", 0]}, + "_meta": {"title": "Modal Artifact Finalizer"}, + } + finalizer_class = fake_nodes_module.NODE_CLASS_MAPPINGS[ + modal_executor_module.MODAL_ARTIFACT_FINALIZER_NODE_ID + ] + assert finalizer_class.GET_SCHEMA().is_output_node is True + assert finalizer_class.OUTPUT_NODE is True + + def test_rewrite_strips_prompt_id_from_cache_safe_proxy_payload( api_intercept_module: Any, settings_module: Any, @@ -1162,7 +1264,7 @@ def test_rewrite_colocates_remote_chain_across_large_transportable_edges( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "3"} + assert set(rewritten_prompt) == {"1", "3", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} @@ -1263,7 +1365,7 @@ def test_rewrite_colocates_sampler_decoders_and_video_assembly( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "8"} + assert set(rewritten_prompt) == {"1", "8", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == { "1": ["1", "2", "3", "4", "5", "6", "7"], @@ -1367,7 +1469,12 @@ def test_rewrite_keeps_non_returning_local_preview_taps_local( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "3", "9"} + assert set(rewritten_prompt) == { + "1", + "3", + "9", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} @@ -1515,7 +1622,17 @@ def test_rewrite_keeps_unmarked_preview_subgraph_nodes_local( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "3", "7", "8", "9", "11", "90", "192"} + assert set(rewritten_prompt) == { + "1", + "3", + "7", + "8", + "9", + "11", + "90", + "192", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} @@ -1613,7 +1730,13 @@ def test_rewrite_keeps_local_branches_that_feed_remote_as_boundaries( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "2", "3", "4"} + assert set(rewritten_prompt) == { + "1", + "2", + "3", + "4", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["1", "2"] assert summary.component_node_ids_by_representative == {"1": ["1"], "2": ["2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "2"} @@ -2015,7 +2138,7 @@ def test_rewrite_merges_cyclic_coarse_components_back_into_single_proxy( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "4"} + assert set(rewritten_prompt) == {"1", "4", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2", "3"]} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] @@ -2118,7 +2241,12 @@ def test_rewrite_marks_modal_map_boundary_as_mapped_subgraph( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "2", "4"} + assert set(rewritten_prompt) == { + "1", + "2", + "4", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["2"] payload = rewritten_prompt["2"]["inputs"]["original_node_data"] assert payload["payload_kind"] == "mapped_subgraph" @@ -2251,7 +2379,13 @@ def test_rewrite_marks_local_modal_map_source_as_mapped_subgraph( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "2", "3", "4"} + assert set(rewritten_prompt) == { + "1", + "2", + "3", + "4", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["3"] payload = rewritten_prompt["3"]["inputs"]["original_node_data"] assert payload["payload_kind"] == "mapped_subgraph" @@ -2403,7 +2537,15 @@ def test_rewrite_supports_mapped_branch_that_shares_non_transportable_upstream_w nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "1__mapped", "2", "4", "5", "8"} + assert set(rewritten_prompt) == { + "1", + "1__mapped", + "2", + "4", + "5", + "8", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["1", "1__mapped"] assert summary.component_node_ids_by_representative == { "1": ["1", "3"], @@ -2710,7 +2852,14 @@ def test_rewrite_keeps_unmapped_remote_siblings_without_local_reentry_together( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "2", "4", "5", "7"} + assert set(rewritten_prompt) == { + "1", + "2", + "4", + "5", + "7", + _artifact_finalizer_node_id(summary), + } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == { "1": ["1", "3", "6"], @@ -3079,7 +3228,7 @@ def test_rewrite_rejects_non_transportable_remote_inputs( nodes_module=fake_nodes_module, ) - assert list(rewritten_prompt) == ["1"] + assert list(rewritten_prompt) == ["1", _artifact_finalizer_node_id(summary)] assert summary.remote_node_ids == ["1", "2"] assert summary.remote_component_ids == ["1"] @@ -3161,7 +3310,7 @@ def test_rewrite_detects_remote_marker_inside_nested_subgraph_workflow( nodes_module=fake_nodes_module, ) - assert list(rewritten_prompt) == ["99", "4"] + assert list(rewritten_prompt) == ["99", "4", _artifact_finalizer_node_id(summary)] assert rewritten_prompt["4"]["inputs"]["latent"] == ["99", 0] assert summary.remote_node_ids == ["99"] assert summary.remote_component_ids == ["99"] @@ -3243,7 +3392,7 @@ def test_rewrite_detects_marked_inner_subgraph_prompt_node_ids( nodes_module=fake_nodes_module, ) - assert list(rewritten_prompt) == ["24:23"] + assert list(rewritten_prompt) == ["24:23", _artifact_finalizer_node_id(summary)] assert summary.remote_node_ids == ["24:23", "30"] assert summary.remote_component_ids == ["24:23"] @@ -3333,7 +3482,7 @@ def test_rewrite_auto_expands_upstream_non_transportable_dependencies( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"1", "4"} + assert set(rewritten_prompt) == {"1", "4", _artifact_finalizer_node_id(summary)} assert summary.remote_node_ids == ["1", "2", "3"] assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2", "3"]} @@ -3857,7 +4006,12 @@ def test_rewrite_keeps_nested_remote_nodes_remote_when_root_ids_collide( nodes_module=fake_nodes_module, ) - assert set(rewritten_prompt) == {"27", "195:27", "223"} + assert set(rewritten_prompt) == { + "27", + "195:27", + "223", + _artifact_finalizer_node_id(summary), + } assert summary.remote_node_ids == ["195:27", "222"] assert summary.remote_component_ids == ["195:27"] diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 25023bc..ec0edb2 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -601,6 +601,78 @@ def test_dynamic_proxy_node_preserves_output_signature( assert proxy_class.INPUT_IS_LIST is True +def test_component_proxy_emits_completion_token_after_remote_execution( + modal_executor_module: Any, +) -> None: + """Artifact-only component proxies should expose completion after remote work returns.""" + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={}, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered( + output_types=(), + output_names=(), + output_is_list=(), + nodes_module=fake_nodes_module, + is_output_node=False, + include_completion_output=True, + ) + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] + schema = proxy_class.GET_SCHEMA() + + assert [output.io_type for output in schema.outputs] == ["BOOLEAN"] + assert [output.display_name for output in schema.outputs] == [ + modal_executor_module.MODAL_COMPONENT_COMPLETION_OUTPUT_NAME + ] + + class FakeClient: + """Return no remote boundary values for an artifact-only component.""" + + async def execute_payload_async( + self, + payload: dict[str, Any], + kwargs: dict[str, Any], + ) -> tuple[()]: + """Confirm the remote component ran before returning no boundary outputs.""" + assert payload["component_id"] == "artifact-component" + assert kwargs == {} + return () + + modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) + try: + result = asyncio.run( + proxy_class.execute( + original_node_data={ + "payload_kind": "subgraph", + "prompt_id": "artifact-prompt", + "component_id": "artifact-component", + }, + unique_id="artifact-component", + ) + ) + finally: + modal_executor_module.set_remote_executor_client_factory(None) + + assert result.result == (True,) + + +def test_modal_artifact_finalizer_requires_completed_components( + modal_executor_module: Any, +) -> None: + """The internal output sink should reject any missing completion signal.""" + schema = modal_executor_module.ModalArtifactFinalizer.GET_SCHEMA() + + assert schema.is_output_node is True + assert schema.outputs == [] + assert modal_executor_module.ModalArtifactFinalizer.execute( + {"component_0": True, "component_1": True} + ).result is None + with pytest.raises(RuntimeError, match="component_1"): + modal_executor_module.ModalArtifactFinalizer.execute( + {"component_0": True, "component_1": False} + ) + + def test_proxy_execution_uses_injected_remote_client( modal_executor_module: Any, ) -> None: @@ -8755,7 +8827,8 @@ def test_split_hybrid_proxies_allow_local_downstream_work_before_mapped_completi assert static_outputs[0] == "static-latent" assert session_state_module.is_remote_session_bridge_ref_payload(static_outputs[1]) - assert mapped_outputs == ("mapped-latent",) + assert static_outputs[2] is True + assert mapped_outputs == ("mapped-latent", True) assert observed_order == [ "static_proxy_finish", "mapped_proxy_start",