"""Tests split from the Modal executor integration suite.""" from __future__ import annotations from modal_executor_test_support import * # noqa: F401,F403 def test_modal_proxy_waits_for_active_prompt_before_starting_next_prompt( modal_executor_module: Any, ) -> None: """A Modal proxy from a later prompt should wait until the active prompt's remote work drains.""" 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=("STRING",), output_names=("value",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] first_started = asyncio.Event() second_started = asyncio.Event() release_first = asyncio.Event() observed_events: list[str] = [] class FakeClient: """Fake async remote client that records when each prompt reaches execution.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str]: """Block the first prompt so the second prompt has to wait at the gate.""" component_id = str(payload["component_id"]) observed_events.append(f"start:{component_id}") if component_id == "component-1": first_started.set() await release_first.wait() if component_id == "component-2": second_started.set() observed_events.append(f"finish:{component_id}") return (component_id,) async def run_scenario() -> tuple[Any, Any]: """Run two prompts concurrently and verify prompt-level serialization.""" first_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-1", }, unique_id="component-1", ) ) await first_started.wait() second_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-2", "component_id": "component-2", }, unique_id="component-2", ) ) await asyncio.sleep(0.05) assert not second_started.is_set() release_first.set() return await asyncio.gather(first_task, second_task) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: first_output, second_output = asyncio.run(run_scenario()) finally: modal_executor_module.set_remote_executor_client_factory(None) assert first_output.result == ("component-1",) assert second_output.result == ("component-2",) assert observed_events == [ "start:component-1", "finish:component-1", "start:component-2", "finish:component-2", ] def test_modal_proxy_allows_same_prompt_components_to_overlap( modal_executor_module: Any, ) -> None: """Modal proxies from the same prompt should be allowed to execute concurrently.""" 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=("STRING",), output_names=("value",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] started_components: set[str] = set() both_started = asyncio.Event() release_components = asyncio.Event() observed_events: list[str] = [] class FakeClient: """Fake async remote client that waits until both same-prompt components start.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str]: """Record concurrent same-prompt starts before returning.""" component_id = str(payload["component_id"]) observed_events.append(f"start:{component_id}") started_components.add(component_id) if started_components == {"component-1", "component-2"}: both_started.set() await release_components.wait() observed_events.append(f"finish:{component_id}") return (component_id,) async def run_scenario() -> tuple[Any, Any, Any]: """Run two same-prompt proxy components concurrently.""" first_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-1", }, unique_id="component-1", ) ) second_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-2", }, unique_id="component-2", ) ) await asyncio.wait_for(both_started.wait(), timeout=1.0) release_components.set() return await asyncio.gather(first_task, second_task) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: first_output, second_output = asyncio.run(run_scenario()) finally: modal_executor_module.set_remote_executor_client_factory(None) assert first_output.result == ("component-1",) assert second_output.result == ("component-2",) assert observed_events[:2] == ["start:component-1", "start:component-2"] def test_modal_proxy_cancellation_while_waiting_does_not_leak_workflow_gate( modal_executor_module: Any, ) -> None: """Cancelling a proxy before it gets the prompt gate should not block later prompts.""" 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=("STRING",), output_names=("value",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] first_started = asyncio.Event() release_first = asyncio.Event() observed_events: list[str] = [] class FakeClient: """Fake async remote client that lets one prompt hold the workflow gate.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str]: """Record remote dispatch and optionally block the first prompt.""" component_id = str(payload["component_id"]) observed_events.append(f"start:{component_id}") if component_id == "component-1": first_started.set() await release_first.wait() observed_events.append(f"finish:{component_id}") return (component_id,) async def run_scenario() -> tuple[Any, Any]: """Cancel a waiting prompt, then prove the next prompt can still dispatch.""" first_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-1", }, unique_id="component-1", ) ) await first_started.wait() waiting_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-2", "component_id": "component-2", }, unique_id="component-2", ) ) await asyncio.sleep(0.05) waiting_task.cancel() with pytest.raises(asyncio.CancelledError): await waiting_task release_first.set() first_output = await first_task third_output = await asyncio.wait_for( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-3", "component_id": "component-3", }, unique_id="component-3", ), timeout=1.0, ) return first_output, third_output modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: first_output, third_output = asyncio.run(run_scenario()) finally: modal_executor_module.set_remote_executor_client_factory(None) assert first_output.result == ("component-1",) assert third_output.result == ("component-3",) assert observed_events == [ "start:component-1", "finish:component-1", "start:component-3", "finish:component-3", ] def test_abandoned_modal_prompt_gate_unblocks_next_prompt( modal_executor_module: Any, ) -> None: """Cancelling an active Modal prompt should not leave later prompts stuck behind its gate.""" 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=("STRING",), output_names=("value",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] started_first_components: set[str] = set() first_components_started = asyncio.Event() second_started = asyncio.Event() release_first_components = asyncio.Event() observed_events: list[str] = [] class FakeClient: """Fake async remote client that lets cancelled prompt components stay stuck.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str]: """Record remote dispatch while prompt one waits for cleanup.""" del kwargs component_id = str(payload["component_id"]) observed_events.append(f"start:{component_id}") if component_id in {"component-1a", "component-1b"}: started_first_components.add(component_id) if started_first_components == {"component-1a", "component-1b"}: first_components_started.set() await release_first_components.wait() if component_id == "component-2": second_started.set() observed_events.append(f"finish:{component_id}") return (component_id,) async def run_scenario() -> tuple[Any, Any]: """Abandon prompt one, then prove prompt two can start before prompt one exits.""" first_task_a = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-1a", }, unique_id="component-1a", ) ) first_task_b = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-1", "component_id": "component-1b", }, unique_id="component-1b", ) ) await first_components_started.wait() second_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": "prompt-2", "component_id": "component-2", }, unique_id="component-2", ) ) await asyncio.sleep(0.05) assert not second_started.is_set() modal_executor_module.abandon_modal_workflow_execution_prompt( "prompt-1", "test cancellation", ) await asyncio.wait_for(second_started.wait(), timeout=1.0) second_output = await second_task release_first_components.set() first_output_a, first_output_b = await asyncio.gather(first_task_a, first_task_b) return first_output_a, first_output_b, second_output modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: first_output_a, first_output_b, second_output = asyncio.run(run_scenario()) finally: modal_executor_module.set_remote_executor_client_factory(None) assert first_output_a.result == ("component-1a",) assert first_output_b.result == ("component-1b",) assert second_output.result == ("component-2",) assert set(observed_events[:2]) == {"start:component-1a", "start:component-1b"} assert observed_events[2:4] == ["start:component-2", "finish:component-2"] assert set(observed_events[4:]) == {"finish:component-1a", "finish:component-1b"} @pytest.mark.parametrize( ("registration_function_name", "node_class_name", "node_id_constant_name"), [ ( "ensure_modal_parallel_local_passthrough_registered", "ModalParallelLocalPassthrough", "MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID", ), ( "ensure_modal_local_bridge_materializer_registered", "ModalLocalBridgeMaterializer", "MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID", ), ( "ensure_modal_artifact_finalizer_registered", "ModalArtifactFinalizer", "MODAL_ARTIFACT_FINALIZER_NODE_ID", ), ], ) def test_internal_node_registration_sets_module_identity( modal_executor_module: Any, monkeypatch: pytest.MonkeyPatch, registration_function_name: str, node_class_name: str, node_id_constant_name: str, ) -> None: """Static internal nodes should serialize the startup node's module identity.""" expected_module = "custom_nodes.ComfyUI-Modal" monkeypatch.setattr( modal_executor_module.ModalUniversalExecutor, "RELATIVE_PYTHON_MODULE", expected_module, ) fake_nodes_module = types.SimpleNamespace( NODE_CLASS_MAPPINGS={}, NODE_DISPLAY_NAME_MAPPINGS={}, ) registration_function = getattr(modal_executor_module, registration_function_name) node_class = getattr(modal_executor_module, node_class_name) monkeypatch.setattr( node_class, "RELATIVE_PYTHON_MODULE", node_class.RELATIVE_PYTHON_MODULE, ) registration_function(fake_nodes_module) node_id = getattr(modal_executor_module, node_id_constant_name) assert fake_nodes_module.NODE_CLASS_MAPPINGS[node_id] is node_class _assert_node_module_identity(node_class, expected_module) def test_dynamic_proxy_node_preserves_output_signature( modal_executor_module: Any, ) -> None: """Dynamic Modal proxies should mirror the original output count and names.""" fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( original_class_type="OriginalNode", original_class=_FakeOriginalNode, nodes_module=fake_nodes_module, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] schema = proxy_class.GET_SCHEMA() assert schema.node_id == proxy_id assert [output.display_name for output in schema.outputs] == ["image", "count"] assert [output.io_type for output in schema.outputs] == ["IMAGE", "INT"] assert proxy_class.INPUT_IS_LIST is True def test_scheduler_list_proxy_wraps_singleton_image_tensor( modal_executor_module: Any, ) -> None: """A singleton mapped IMAGE must retain its batch dimension in ComfyUI.""" torch = pytest.importorskip("torch") 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=("IMAGE",), output_names=("image",), output_is_list=(True,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] image = torch.zeros((1, 8, 6, 3), dtype=torch.float32) class FakeClient: """Return one ordinary IMAGE tensor from a singleton remote execution.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[Any, ...]: """Return the unwrapped tensor produced by the ordinary subgraph path.""" assert payload["boundary_outputs"][0]["scheduler_is_list"] is True assert kwargs == {} return (image,) 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": "singleton-image-prompt", "component_id": "singleton-image-component", "boundary_outputs": [ { "proxy_output_name": "image", "node_id": "11", "output_index": 0, "io_type": "IMAGE", "is_list": False, "scheduler_is_list": True, } ], }, unique_id="singleton-image-component", ) ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert isinstance(result.result[0], list) assert len(result.result[0]) == 1 assert result.result[0][0] is image scheduler_items: list[Any] = [] scheduler_items.extend(result.result[0]) assert len(scheduler_items) == 1 assert scheduler_items[0] is image assert tuple(scheduler_items[0].shape) == (1, 8, 6, 3) def test_transportable_singleton_list_reaches_scalar_remote_consumer_as_item( modal_executor_module: Any, ) -> None: """A cached singleton list output should enter the next remote proxy as a scalar.""" fake_nodes_module = types.SimpleNamespace( NODE_CLASS_MAPPINGS={}, NODE_DISPLAY_NAME_MAPPINGS={}, ) producer_proxy_id = ( modal_executor_module.ensure_modal_component_proxy_node_registered( output_types=("INT",), output_names=("seed",), output_is_list=(True,), nodes_module=fake_nodes_module, is_output_node=False, ) ) consumer_proxy_id = ( modal_executor_module.ensure_modal_component_proxy_node_registered( output_types=("STRING",), output_names=("result",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) ) producer_proxy = fake_nodes_module.NODE_CLASS_MAPPINGS[producer_proxy_id] consumer_proxy = fake_nodes_module.NODE_CLASS_MAPPINGS[consumer_proxy_id] observed_consumer_inputs: list[Any] = [] class FakeClient: """Return a singleton seed list and record the downstream normalized input.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[Any, ...]: """Model the two remote components around ComfyUI's output cache.""" if payload["component_id"] == "producer": assert kwargs == {} return ([123],) observed_consumer_inputs.append(kwargs["remote_input_0"]) return (f"seed:{int(kwargs['remote_input_0'])}",) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: producer_result = asyncio.run( producer_proxy.execute( original_node_data={ "payload_kind": "subgraph", "component_id": "producer", "boundary_outputs": [ { "proxy_output_name": "seed", "node_id": "325", "output_index": 0, "io_type": "INT", "is_list": True, } ], }, unique_id="producer", ) ) cached_scheduler_items: list[Any] = [] cached_scheduler_items.extend(producer_result.result[0]) consumer_result = asyncio.run( consumer_proxy.execute( original_node_data={ "payload_kind": "subgraph", "component_id": "consumer", "boundary_outputs": [ { "proxy_output_name": "result", "node_id": "372", "output_index": 0, "io_type": "STRING", "is_list": False, } ], }, remote_input_0=cached_scheduler_items, unique_id="consumer", ) ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert cached_scheduler_items == [123] assert observed_consumer_inputs == [123] assert consumer_result.result == ("seed:123",) 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_modal_parallel_local_passthrough_releases_after_remote_dispatch( modal_executor_module: Any, ) -> None: """A local preview should unblock when its remote continuation starts.""" schema = modal_executor_module.ModalParallelLocalPassthrough.GET_SCHEMA() dispatch_context = { "dispatch_group_id": "prompt-parallel", "component_ids": ["component-b"], } async def run_scenario() -> tuple[Any, ...]: """Prove the passthrough remains pending until remote dispatch is signalled.""" passthrough_task = asyncio.create_task( modal_executor_module.ModalParallelLocalPassthrough.execute( "image-value", dispatch_context, ) ) await asyncio.sleep(0) assert passthrough_task.done() is False assert modal_executor_module._signal_parallel_local_dispatch( { "parallel_local_dispatch_group_id": "prompt-parallel", "component_id": "component-b", "signal_parallel_local_dispatch": True, } ) is True result = await asyncio.wait_for(passthrough_task, timeout=1.0) return result.result try: result = asyncio.run(run_scenario()) finally: modal_executor_module._clear_parallel_dispatch_events("prompt-parallel") assert schema.is_output_node is False assert result == ("image-value",) def test_modal_parallel_local_passthrough_runs_before_remote_result_returns( modal_executor_module: Any, ) -> None: """A local preview should run while its downstream remote call is in flight.""" 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=("STRING",), output_names=("value",), output_is_list=(False,), nodes_module=fake_nodes_module, is_output_node=False, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] remote_started = asyncio.Event() release_remote = asyncio.Event() class FakeClient: """Fake async remote client that remains in flight until explicitly released.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str]: """Record dispatch, then keep the remote result pending.""" remote_started.set() await release_remote.wait() return (str(payload["component_id"]),) async def run_scenario() -> tuple[tuple[Any, ...], tuple[Any, ...]]: """Prove local passthrough completion precedes the remote result.""" dispatch_group_id = "prompt-parallel-in-flight" local_task = asyncio.create_task( modal_executor_module.ModalParallelLocalPassthrough.execute( "preview-image", { "dispatch_group_id": dispatch_group_id, "component_ids": ["component-b"], }, ) ) remote_task = asyncio.create_task( proxy_class.execute( original_node_data={ "payload_kind": "subgraph", "prompt_id": dispatch_group_id, "component_id": "component-b", "parallel_local_dispatch_group_id": dispatch_group_id, "signal_parallel_local_dispatch": True, }, unique_id="component-b", ) ) await asyncio.wait_for(remote_started.wait(), timeout=1.0) local_result = await asyncio.wait_for(local_task, timeout=1.0) assert remote_task.done() is False release_remote.set() remote_result = await asyncio.wait_for(remote_task, timeout=1.0) return local_result.result, remote_result.result modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: local_result, remote_result = asyncio.run(run_scenario()) finally: modal_executor_module.set_remote_executor_client_factory(None) modal_executor_module._clear_parallel_dispatch_events( "prompt-parallel-in-flight" ) assert local_result == ("preview-image",) assert remote_result == ("component-b",) def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component( api_intercept_module: Any, prompt_affinity_planning_module: Any, remote_graph_analysis_module: Any, ) -> None: """A preview gate should not wait for remote descendants of a remote frontier.""" fake_nodes_module = types.SimpleNamespace( NODE_CLASS_MAPPINGS={ "RemoteProxy": _FakeRewriteRemoteSamplerNode, "LocalFeedback": _FakeRewriteLocalFeedbackNode, }, ) rewritten_prompt = { "source": {"class_type": "RemoteProxy", "inputs": {}}, "feedback": { "class_type": "LocalFeedback", "inputs": {"value": ["source", 0]}, }, "nearest": { "class_type": "RemoteProxy", "inputs": {"value": ["feedback", 0]}, }, "later": { "class_type": "RemoteProxy", "inputs": {"value": ["nearest", 0]}, }, } consumers = remote_graph_analysis_module._build_consumer_map(rewritten_prompt) nearest_component_ids = ( prompt_affinity_planning_module._nearest_downstream_remote_component_ids( rewritten_prompt=rewritten_prompt, consumers=consumers, seed_targets=consumers[ api_intercept_module.LinkedOutputRef("source", 0) ], remote_component_id_set={"source", "nearest", "later"}, nodes_module=fake_nodes_module, ) ) assert nearest_component_ids == ["nearest"] def test_proxy_execution_uses_injected_remote_client( modal_executor_module: Any, ) -> None: """Proxy execution should delegate to the configured remote client asynchronously.""" fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( original_class_type="OriginalNode", original_class=_FakeOriginalNode, nodes_module=fake_nodes_module, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] class FakeClient: """Test client that returns deterministic outputs.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str, int]: """Return values derived from the proxied node payload.""" return (f"{payload['class_type']}::{kwargs['value']}", 3) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: result = asyncio.run( proxy_class.execute(original_node_data={"class_type": "OriginalNode"}, value="payload") ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert result.result == ("OriginalNode::payload", 3) def test_proxy_execution_normalizes_input_is_list_kwargs( modal_executor_module: Any, ) -> None: """Dynamic Modal proxies should unwrap singleton INPUT_IS_LIST wrappers before remote execution.""" fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( original_class_type="OriginalNode", original_class=_FakeOriginalNode, nodes_module=fake_nodes_module, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] observed_kwargs: dict[str, Any] = {} class FakeClient: """Test client that records normalized proxy kwargs.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str, int]: """Capture the kwargs forwarded by the proxy.""" observed_kwargs["payload_kind"] = payload["payload_kind"] observed_kwargs.update(kwargs) return ("ok", 1) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: result = asyncio.run( proxy_class.execute( original_node_data=[{"payload_kind": "subgraph", "class_type": "OriginalNode"}], scalar_value=[3], mapped_value=["a", "b", "c"], ) ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert result.result == ("ok", 1) assert observed_kwargs == { "payload_kind": "subgraph", "scalar_value": 3, "mapped_value": ["a", "b", "c"], } def test_proxy_execution_wraps_sync_remote_clients( modal_executor_module: Any, ) -> None: """Async proxy execution should still support legacy sync remote clients.""" fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( original_class_type="OriginalNode", original_class=_FakeOriginalNode, nodes_module=fake_nodes_module, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] class FakeSyncClient: """Legacy client that only exposes the blocking execution method.""" def execute_payload(self, payload: dict[str, Any], kwargs: dict[str, Any]) -> tuple[str, int]: """Return values derived from the proxied node payload.""" return (f"sync::{payload['class_type']}::{kwargs['value']}", 4) modal_executor_module.set_remote_executor_client_factory(lambda: FakeSyncClient()) try: result = asyncio.run( proxy_class.execute(original_node_data={"class_type": "OriginalNode"}, value="payload") ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert result.result == ("sync::OriginalNode::payload", 4) def test_local_remote_app_executes_original_node( remote_modal_app_module: Any, serialization_module: Any, ) -> None: """The local fallback remote app should execute a mapped legacy node.""" payload = remote_modal_app_module.execute_node_locally( node_data={"class_type": "OriginalNode"}, kwargs_payload='{"value": "hello"}', node_mapping={"OriginalNode": _FakeOriginalNode}, ) outputs = serialization_module.deserialize_node_outputs(payload) assert outputs == ("hello", 1) def test_local_gap_keepalive_is_bounded_and_stopped_by_next_remote_component( remote_modal_app_module: Any, modal_warmup_module: Any, monkeypatch: Any, ) -> None: """A completed producer should retain its slot only until downstream remote work starts.""" submitted_tasks: list[tuple[Any, tuple[Any, ...], Future[Any]]] = [] class FakeExecutor: """Executor double that captures but does not run the keepalive loop.""" def submit(self, fn: Any, *args: Any) -> Future[Any]: """Record one keepalive job and return its pending future.""" future: Future[Any] = Future() submitted_tasks.append((fn, args, future)) return future monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote") monkeypatch.setattr(remote_modal_app_module, "modal", object()) monkeypatch.setattr(modal_warmup_module, "modal", object()) monkeypatch.setattr( modal_warmup_module, "_REMOTE_MODAL_KEEPALIVE_EXECUTOR", FakeExecutor(), ) remote_modal_app_module.get_settings.cache_clear() with modal_warmup_module._LOCAL_GAP_KEEPALIVES_LOCK: modal_warmup_module._LOCAL_GAP_KEEPALIVES.clear() producer_payload = { "prompt_id": "prompt-gap", "component_id": "component-a", "execution_provider": "modal", "execution_environment_id": "modal:RTX-PRO-6000", "keepalive_after_remote_component": True, } consumer_payload = { "prompt_id": "prompt-gap", "component_id": "component-b", "execution_provider": "modal", "execution_environment_id": "modal:RTX-PRO-6000", "stop_local_gap_keepalive_before_remote_component": True, } try: assert modal_warmup_module._start_local_gap_keepalive(producer_payload) is True assert len(submitted_tasks) == 1 _fn, args, _future = submitted_tasks[0] stop_event = args[1] assert args[2:] == (900.0, 15.0) assert stop_event.is_set() is False assert modal_warmup_module._stop_local_gap_keepalive( consumer_payload, reason="test_next_component", ) is True assert stop_event.is_set() is True finally: remote_modal_app_module.get_settings.cache_clear() with modal_warmup_module._LOCAL_GAP_KEEPALIVES_LOCK: modal_warmup_module._LOCAL_GAP_KEEPALIVES.clear() def test_register_exact_component_parallelism_refines_prompt_target( remote_modal_app_module: Any, modal_warmup_module: Any, monkeypatch: Any, ) -> None: """Mapped fan-out should raise the prompt-wide warmup target once exact item count is known.""" monkeypatch.setenv("COMFY_MODAL_MAX_CONTAINERS", "6") remote_modal_app_module.get_settings.cache_clear() with modal_warmup_module._PROMPT_WARMUP_STATES_LOCK: modal_warmup_module._PROMPT_WARMUP_STATES.clear() modal_warmup_module._PROMPT_WARMUP_STATE_ORDER = None try: payload = { "prompt_id": "prompt-2", "component_id": "component-a", "extra_data": { "modal": { "component_execution_stages": [["component-a", "component-b"], ["component-c"]], "mapped_component_ids": ["component-a"], "estimated_max_parallel_requests": 2, } }, } refined_target = modal_warmup_module._register_exact_component_parallelism(payload, 5) finally: remote_modal_app_module.get_settings.cache_clear() assert refined_target == 6 def test_workflow_gpu_changes_expected_remote_runtime_fingerprint( remote_modal_app_module: Any, modal_deployment_module: Any, ) -> None: """Each workflow GPU target should select a distinct deploy-time runtime identity.""" a100_fingerprint = modal_deployment_module._expected_remote_runtime_fingerprint( {"modal_gpu": "A100"} ) b300_fingerprint = modal_deployment_module._expected_remote_runtime_fingerprint( {"modal_gpu": "B300"} ) assert a100_fingerprint != b300_fingerprint assert modal_deployment_module._settings_for_payload( {"modal_gpu": "L40S"} ).modal_gpu == "L40S" configured_settings = modal_deployment_module._settings_for_payload( {"modal_gpu": "H200", "modal_max_containers": 3} ) assert configured_settings.modal_gpu == "H200" assert configured_settings.max_containers == 3 assert modal_deployment_module._modal_deploy_cache_key( {"modal_gpu": "A100"} )[0] == "comfy-modal-sync" assert modal_deployment_module._modal_deploy_cache_key( {"modal_gpu": "B300"} )[0] == "comfy-modal-sync-gpu-b300" def test_local_gap_components_share_one_affinity_slot( remote_modal_app_module: Any, modal_deployment_module: Any, ) -> None: """Sequential remote phases around local work must address the same worker pool.""" parallelism_metadata = { "modal": { "component_execution_stages": [ ["component-a", "independent-component"], ["component-b"], ], "estimated_max_parallel_requests": 2, } } first_payload = { "component_id": "component-a", "remote_local_gap_pool": True, "extra_data": parallelism_metadata, } second_payload = { "component_id": "component-b", "remote_local_gap_pool": True, "extra_data": parallelism_metadata, } assert modal_deployment_module._component_pool_slot_index(first_payload) == 0 assert modal_deployment_module._component_pool_slot_index(second_payload) == 0 assert modal_deployment_module._remote_worker_affinity_key(first_payload) == ( modal_deployment_module._remote_worker_affinity_key(second_payload) ) def test_emit_local_remote_dispatch_status_marks_component_starting( remote_modal_app_module: Any, local_ui_events_module: Any, monkeypatch: Any, ) -> None: """Dispatching a remote component should immediately tell the local UI it is starting.""" class FakePromptServer: """Capture websocket events emitted by local dispatch status.""" def __init__(self) -> None: """Initialize the event sink.""" self.messages: list[tuple[str, dict[str, Any], str | None]] = [] def send_sync(self, event: str, data: dict[str, Any], sid: str | None) -> None: """Record one emitted websocket message.""" self.messages.append((event, data, sid)) prompt_server = FakePromptServer() monkeypatch.setattr(local_ui_events_module, "_lookup_local_prompt_server", lambda: prompt_server) remote_modal_app_module._emit_local_remote_dispatch_status( { "prompt_id": "prompt-1", "component_id": "component-1", "component_node_ids": ["7", "8"], "modal_gpu": "B300", "extra_data": {"client_id": "client-1"}, } ) assert prompt_server.messages == [ ( "modal_status", { "phase": "starting", "prompt_id": "prompt-1", "node_ids": ["7", "8"], "modal_gpu": "B300", "status_message": "Starting remote component on Modal", }, "client-1", ) ] def test_modal_gpu_estimated_rates_cover_supported_aliases( modal_container_logs_module: Any, ) -> None: """Published GPU estimates should cover every selectable billed GPU identity.""" rates = modal_container_logs_module.MODAL_GPU_ESTIMATED_USD_PER_SECOND assert set(rates) == set(modal_container_logs_module.MODAL_GPU_TYPES) assert rates["A100"] == rates["A100-40GB"] assert rates["H100!"] == rates["H100"] assert rates["B200+"] == rates["B200"] assert modal_container_logs_module.MODAL_GPU_PRICING_EFFECTIVE_DATE == "2026-08-13" def test_request_remote_interrupt_async_uses_shared_control_async_put( remote_modal_app_module: Any, modal_interrupts_module: Any, monkeypatch: Any, ) -> None: """Async local cancellation should use Modal Dict.put.aio instead of blocking put.""" class FakePut: """Modal Dict put double that records sync and async writes separately.""" def __init__(self) -> None: """Initialize captured writes.""" self.sync_calls: list[tuple[str, Any]] = [] self.async_calls: list[tuple[str, Any]] = [] def __call__(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool: """Record one blocking write.""" del skip_if_exists self.sync_calls.append((key, value)) return True async def aio(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool: """Record one async write.""" del skip_if_exists self.async_calls.append((key, value)) return True class FakeInterruptStore: """Simple Modal Dict double exposing put.aio.""" def __init__(self) -> None: """Initialize the fake async put handle.""" self.put = FakePut() interrupt_store = FakeInterruptStore() class FakeModalDict: """Minimal modal.Dict shim that returns the fake interrupt store.""" @staticmethod def from_name( name: str, *, environment_name: str | None = None, create_if_missing: bool = False, client: Any | None = None, ) -> FakeInterruptStore: return interrupt_store monkeypatch.setattr( modal_interrupts_module, "modal", types.SimpleNamespace(Dict=FakeModalDict), ) monkeypatch.setenv("COMFY_MODAL_INTERRUPT_DICT_NAME", "shared-interrupts") remote_modal_app_module.get_settings.cache_clear() modal_interrupts_module._MODAL_INTERRUPT_DICTS.clear() try: wrote_interrupt = asyncio.run( remote_modal_app_module._request_remote_interrupt_async( {"prompt_id": "prompt-1", "component_id": "component-2"} ) ) finally: remote_modal_app_module.get_settings.cache_clear() modal_interrupts_module._MODAL_INTERRUPT_DICTS.clear() assert wrote_interrupt is True assert interrupt_store.put.sync_calls == [] assert len(interrupt_store.put.async_calls) == 1 interrupt_key, interrupt_value = interrupt_store.put.async_calls[0] assert interrupt_key == "prompt-1:component-2" assert isinstance(interrupt_value["requested_at"], float) def test_local_remote_app_normalizes_wrapped_scalar_prompt_inputs( remote_modal_app_module: Any, serialization_module: Any, ) -> None: """The local fallback runner should unwrap singleton-list scalar prompt inputs.""" payload = remote_modal_app_module.execute_subgraph_locally( payload={ "payload_kind": "subgraph", "component_id": "component-1", "subgraph_prompt": { "remote_1": { "class_type": "BoundarySource", "inputs": {"value": [4]}, "_meta": {}, } }, "boundary_inputs": [], "boundary_outputs": [ { "proxy_output_name": "remote_1_value", "node_id": "remote_1", "output_index": 0, "io_type": "INT", "is_list": False, } ], "execute_node_ids": [["remote_1"]], "extra_data": {}, "custom_nodes_bundle": None, }, kwargs_payload="{}", node_mapping={ "BoundarySource": _BoundarySourceNode, }, ) outputs = serialization_module.deserialize_node_outputs(payload) assert outputs == (5,) @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_apply_boundary_inputs_normalizes_wrapped_scalar_values( request: Any, module_fixture_name: str, ) -> None: """Boundary input hydration should unwrap singleton-list scalar wrappers before PromptExecutor sees them.""" target_module = request.getfixturevalue(module_fixture_name) prompt = { "remote_1": { "class_type": "BoundarySource", "inputs": {"value": 0}, "_meta": {}, } } target_module._apply_boundary_inputs( prompt=prompt, boundary_input_specs=[ { "proxy_input_name": "remote_input_0", "targets": [{"node_id": "remote_1", "input_name": "value"}], } ], hydrated_inputs={"remote_input_0": [4]}, ) assert prompt["remote_1"]["inputs"]["value"] == 4 @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_apply_boundary_inputs_preserves_singleton_conditioning_lists( request: Any, module_fixture_name: str, ) -> None: """Boundary input hydration must not flatten singleton CONDITIONING payloads.""" target_module = request.getfixturevalue(module_fixture_name) conditioning = [["cond", {"pooled_output": None}]] prompt = { "remote_1": { "class_type": "BoundarySource", "inputs": {"value": 0}, "_meta": {}, } } target_module._apply_boundary_inputs( prompt=prompt, boundary_input_specs=[ { "proxy_input_name": "remote_input_0", "io_type": "CONDITIONING", "targets": [{"node_id": "remote_1", "input_name": "value"}], } ], hydrated_inputs={"remote_input_0": conditioning}, ) assert prompt["remote_1"]["inputs"]["value"] == conditioning @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_validate_required_prompt_inputs_reports_missing_latent_image( request: Any, module_fixture_name: str, ) -> None: """Remote execution should fail before PromptExecutor calls samplers with missing required inputs.""" target_module = request.getfixturevalue(module_fixture_name) prompt = { "3": { "class_type": "KSamplerLoraSigmaInverse", "inputs": { "seed": 1, "steps": 24, "positive": [["cond", {"reference_latents": ["reference"]}]], }, } } with pytest.raises(target_module.RemoteSubgraphExecutionError) as exc_info: target_module._validate_required_prompt_inputs( prompt, {"KSamplerLoraSigmaInverse": _FakeImplicitBatchKSamplerNode}, ) message = str(exc_info.value) assert "missing required node inputs" in message assert "latent_image" in message assert "available_inputs=['positive', 'seed', 'steps']" in message @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_validate_required_prompt_inputs_expands_v3_autogrow_inputs( request: Any, module_fixture_name: str, monkeypatch: pytest.MonkeyPatch, ) -> None: """V3 Autogrow prompts should validate their expanded `group.name` sockets.""" target_module = request.getfixturevalue(module_fixture_name) _install_fake_v3_input_finalizer(monkeypatch) batch_images_node = _v3_batch_images_node_class() live_inputs = { "images.image0": ["101", 0], "images.image1": ["102", 0], } prompt = { "153": { "class_type": "ModalTestV3BatchImagesNode", "inputs": live_inputs, } } target_module._validate_required_prompt_inputs( prompt, {"ModalTestV3BatchImagesNode": batch_images_node}, ) assert target_module._node_required_input_names( batch_images_node, live_inputs, ) == {"images.image0"} input_type_map = target_module._node_input_type_map(batch_images_node, live_inputs) assert input_type_map["images.image0"] == "IMAGE" assert input_type_map["images.image1"] == "IMAGE" assert "images" not in input_type_map @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_validate_prompt_input_shapes_rejects_list_on_primitive_socket( request: Any, module_fixture_name: str, ) -> None: """Prepared remote prompts should fail early when primitive widget inputs still carry raw lists.""" target_module = request.getfixturevalue(module_fixture_name) prompt = { "remote_1": { "class_type": "BoundarySource", "inputs": {"value": [4, 5]}, "_meta": {}, } } with pytest.raises(target_module.RemoteSubgraphExecutionError, match="input_name='value'"): target_module._validate_prompt_input_shapes( prompt, {"BoundarySource": _BoundarySourceNode}, ) @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_validate_prompt_input_shapes_allows_boundary_supplied_list_on_primitive_socket( request: Any, module_fixture_name: str, ) -> None: """Boundary-fed primitive lists should pass through so ComfyUI can map over them normally.""" target_module = request.getfixturevalue(module_fixture_name) prompt = { "remote_1": { "class_type": "BoundarySource", "inputs": {"value": [4, 5]}, "_meta": {}, } } target_module._validate_prompt_input_shapes( prompt, {"BoundarySource": _BoundarySourceNode}, [ { "proxy_input_name": "remote_input_0", "targets": [{"node_id": "remote_1", "input_name": "value"}], } ], ) @pytest.mark.parametrize( ("module_fixture_name",), [ ("local_execution_module",), ("modal_cloud_module",), ], ) def test_coerce_prompt_primitive_input_values_matches_comfyui_semantics( request: Any, module_fixture_name: str, ) -> None: """Remote runtimes should coerce primitive prompt literals the same way ComfyUI does.""" target_module = request.getfixturevalue(module_fixture_name) prompt = { "remote_1": { "class_type": "PrimitiveEcho", "inputs": { "steps": 18.0, "cfg": 5, "enabled": 1, "label": 7, }, "_meta": {}, } } target_module._coerce_prompt_primitive_input_values( prompt, {"PrimitiveEcho": _PrimitiveEchoNode}, ) assert prompt["remote_1"]["inputs"] == { "steps": 18, "cfg": 5.0, "enabled": True, "label": "7", } assert isinstance(prompt["remote_1"]["inputs"]["steps"], int) assert isinstance(prompt["remote_1"]["inputs"]["cfg"], float) assert isinstance(prompt["remote_1"]["inputs"]["enabled"], bool) assert isinstance(prompt["remote_1"]["inputs"]["label"], str) def test_local_remote_app_coerces_primitive_widget_literals_before_execution( remote_modal_app_module: Any, serialization_module: Any, ) -> None: """The local fallback runner should coerce primitive widget literals before executing nodes.""" payload = remote_modal_app_module.execute_subgraph_locally( payload={ "payload_kind": "subgraph", "component_id": "component-primitive-coercion", "subgraph_prompt": { "remote_1": { "class_type": "PrimitiveEcho", "inputs": { "steps": 18.0, "cfg": 5, "enabled": 1, "label": 7, }, "_meta": {}, } }, "boundary_inputs": [], "boundary_outputs": [ { "proxy_output_name": "steps", "node_id": "remote_1", "output_index": 0, "io_type": "INT", "is_list": False, }, { "proxy_output_name": "cfg", "node_id": "remote_1", "output_index": 1, "io_type": "FLOAT", "is_list": False, }, { "proxy_output_name": "enabled", "node_id": "remote_1", "output_index": 2, "io_type": "BOOLEAN", "is_list": False, }, { "proxy_output_name": "label", "node_id": "remote_1", "output_index": 3, "io_type": "STRING", "is_list": False, }, ], "execute_node_ids": ["remote_1"], "extra_data": {}, "custom_nodes_bundle": None, }, kwargs_payload="{}", node_mapping={"PrimitiveEcho": _PrimitiveEchoNode}, ) outputs = serialization_module.deserialize_node_outputs(payload) assert outputs == (18, 5.0, True, "7")