diff --git a/README.md b/README.md index 8f8cadc..97fcd4d 100644 --- a/README.md +++ b/README.md @@ -396,6 +396,7 @@ COMFYUI_ROOT=/tmp/comfyui-modal-test/ComfyUI \ - Prompt-scoped remote session misses are treated as routing/state errors now, not poisoned-container crashes. If a split proxy lands on the wrong Modal container and its session bucket is missing, the worker logs the miss but does not self-terminate; the real fix for that class of failure is stronger container affinity or leased session routing. - `ModalMapInput` is a queue-time mapping marker even when the marker node itself stays local. If a local `ModalMapInput` feeds a remote node, rewrite should still produce a `mapped_subgraph` so list(INT) values like `NextSeeds` outputs are itemized before the remote sampler runs. - Implicit mapped execution distinguishes semantic lists from lists of remote session refs. Raw `CONDITIONING` lists stay broadcast, but a list of prompt-scoped session refs for `MODEL`/`CONDITIONING` is treated as per-item mapped state and split item-by-item before the downstream remote node executes. +- Implicit mapped execution cannot let each static or per-item sub-run inherit `clear_remote_session=True`. Shared prompt-scoped sessions now stay alive across all sibling mapped items and are cleared exactly once through a dedicated final cleanup payload, which prevents later item runs from seeing `RemoteSessionStateError` after an earlier sibling finishes first. ## Current limitations diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 3093b03..79f38cf 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -158,3 +158,4 @@ - Live combined logs confirmed the current split-proxy weak point: `session_affinity_key` is best-effort, not a hard lease. One prompt-scoped session can be created on container `A` and a later phase can still land on container `B`, producing a `RemoteSessionStateError` even though neither worker is corrupted. That class of failure should not trigger container teardown; it needs stronger routing semantics. - `ModalMapInput` cannot be treated as a remote-only marker. If the marker stays local and only the downstream sampler is remote-marked, rewrite still has to recognize that boundary as mapped work. Otherwise the remote component falls back to an ordinary `subgraph`, the boundary io_type degrades to `*`, and batched primitive lists like `list(INT)` seeds stop itemizing before they hit the sampler. - Nontransportable mapped state needs a different batching rule from raw semantic payloads. A plain `CONDITIONING` Python list usually encodes one conditioning structure and must stay broadcast, but a Python list whose elements are all remote session refs encodes per-item mapped remote values. Those ref lists need to split item-by-item even when the declared io_type is `MODEL` or `CONDITIONING`. +- Implicit mapped fan-out also needs its own session-lifetime rule. If every synthesized `::item:*` or `::static` payload inherits `clear_remote_session=True`, the first sibling to finish can tear down shared remote state while later siblings are still queued. Cleanup has to move into one explicit final payload that runs after all mapped siblings complete. diff --git a/docs/UPNEXT.md b/docs/UPNEXT.md index 8100526..d6c2804 100644 --- a/docs/UPNEXT.md +++ b/docs/UPNEXT.md @@ -47,5 +47,6 @@ Partial-completion semantics for one proxy would require ComfyUI scheduler behav - Combined local+remote logs also exposed the real remaining split-proxy risk under load: a later phase can still land on a different Modal container and miss the prompt-scoped session bucket entirely. We now avoid killing that second container on session misses, but the underlying product fix is still stronger container affinity or explicit leased-worker/session routing. - Rewrite also has to honor local `ModalMapInput` markers. That case is now covered so `list(INT)` outputs from local nodes like `NextSeeds` still become mapped remote item runs instead of ordinary remote subgraphs that pass the whole list into one sampler call. - Implicit batching also has to recognize per-item remote session refs for nontransportable types. Lists of `MODEL`/`CONDITIONING` session refs now split item-by-item, while raw conditioning payload lists still stay broadcast. +- Implicit mapped cleanup now also runs through one dedicated final payload instead of letting each `::static` or `::item:*` sub-run inherit `clear_remote_session=True`. That keeps shared prompt-scoped state alive until all sibling mapped runs finish. - Run with `COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS=true` during broader production workflows and watch the mirrored session lifecycle logs for any real session misses under scale-out. - If production logs ever show split-proxy session misses under real scale-out, add a stronger affinity fallback or explicit leased-worker/session routing instead of relying on best-effort class-instance affinity. diff --git a/remote/modal_app.py b/remote/modal_app.py index 8ffb05b..41e1305 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -2300,6 +2300,7 @@ def _build_mapped_item_payload( item_payload["mapped_progress_display_node_id"] = str( payload.get("component_id", "modal-subgraph") ) + item_payload.pop("clear_remote_session", None) item_payload["execute_node_ids"] = list( payload.get("mapped_execute_node_ids") or payload.get("execute_node_ids", []) ) @@ -2337,6 +2338,27 @@ def _aggregate_mapped_outputs( return tuple(aggregated_outputs) +def _build_remote_session_cleanup_payload(payload: dict[str, Any]) -> dict[str, Any] | None: + """Return a dedicated cleanup payload that clears one shared remote session once.""" + remote_session = payload.get("remote_session") + if not bool(payload.get("clear_remote_session")) or not is_remote_session_handle_payload(remote_session): + return None + return { + "payload_kind": "subgraph", + "component_id": f"{payload.get('component_id', 'modal-subgraph')}::cleanup", + **_shared_subgraph_payload_fields(payload), + "component_node_ids": [], + "subgraph_prompt": {}, + "boundary_inputs": [], + "boundary_outputs": [], + "execute_node_ids": [], + "remote_session": copy.deepcopy(remote_session), + "clear_remote_session": True, + "suppress_status_stream": True, + "terminate_container_on_error": False, + } + + def _is_mapped_boundary_output(boundary_output: dict[str, Any], payload: dict[str, Any]) -> bool: """Return whether one boundary output belongs to the mapped per-item branch.""" mapped_output = boundary_output.get("mapped_output") @@ -2359,6 +2381,7 @@ def _build_static_mapped_payload(payload: dict[str, Any]) -> dict[str, Any]: static_payload["component_id"] = f"{payload.get('component_id', 'modal-subgraph')}::static" static_payload["mapped_input"] = None static_payload["suppress_status_stream"] = True + static_payload.pop("clear_remote_session", None) static_payload["execute_node_ids"] = list(payload.get("static_execute_node_ids") or []) static_payload["boundary_outputs"] = [ copy.deepcopy(boundary_output) @@ -2685,73 +2708,81 @@ async def _invoke_implicitly_mapped_subgraph_async(payload: dict[str, Any], kwar payload, mapped_execute_node_ids, ) + cleanup_payload = _build_remote_session_cleanup_payload(hybrid_payload) static_outputs: tuple[Any, ...] = () - if static_execute_node_ids: - static_response = await invoke_remote_engine_async( - _build_static_mapped_payload(hybrid_payload), - serialize_node_inputs(broadcast_inputs), - ) - static_outputs = deserialize_node_outputs(static_response) - - per_item_outputs: list[tuple[Any, ...] | None] = [None] * total_items - completed_items = 0 - item_queue: asyncio.Queue[int | None] = asyncio.Queue() - for item_index in range(total_items): - item_queue.put_nowait(item_index) - for _ in range(parallelism): - item_queue.put_nowait(None) - - async def run_worker(lane_index: int) -> None: - """Execute queued implicit mapped items through one stable local worker lane.""" - nonlocal completed_items - while True: - item_index = await item_queue.get() - if item_index is None: - return - if _local_processing_interrupted(): - _raise_local_interrupt() - try: - item_payload = _build_mapped_item_payload(hybrid_payload, item_index, lane_index) - item_inputs = dict(broadcast_inputs) - for input_name, items in split_inputs.items(): - item_inputs[input_name] = items[item_index] - item_response = await invoke_remote_engine_async( - item_payload, - serialize_node_inputs(item_inputs), - ) - per_item_outputs[item_index] = deserialize_node_outputs(item_response) - completed_items += 1 - _emit_local_mapped_progress(payload, completed_items, total_items) - finally: - _clear_local_mapped_lane_progress(payload, lane_index, item_index) - - tasks = [asyncio.create_task(run_worker(lane_index)) for lane_index in range(parallelism)] try: - await asyncio.gather(*tasks) - except Exception: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - raise + if static_execute_node_ids: + static_response = await invoke_remote_engine_async( + _build_static_mapped_payload(hybrid_payload), + serialize_node_inputs(broadcast_inputs), + ) + static_outputs = deserialize_node_outputs(static_response) - return serialize_node_outputs( - _merge_static_and_mapped_outputs( - static_outputs=static_outputs, - mapped_outputs=_aggregate_mapped_outputs( - [item_outputs for item_outputs in per_item_outputs if item_outputs is not None], - { - **hybrid_payload, - "boundary_outputs": [ - boundary_output - for boundary_output in hybrid_payload.get("boundary_outputs", []) - if _is_mapped_boundary_output(boundary_output, hybrid_payload) - ], - }, - ), - payload=hybrid_payload, + per_item_outputs: list[tuple[Any, ...] | None] = [None] * total_items + completed_items = 0 + item_queue: asyncio.Queue[int | None] = asyncio.Queue() + for item_index in range(total_items): + item_queue.put_nowait(item_index) + for _ in range(parallelism): + item_queue.put_nowait(None) + + async def run_worker(lane_index: int) -> None: + """Execute queued implicit mapped items through one stable local worker lane.""" + nonlocal completed_items + while True: + item_index = await item_queue.get() + if item_index is None: + return + if _local_processing_interrupted(): + _raise_local_interrupt() + try: + item_payload = _build_mapped_item_payload(hybrid_payload, item_index, lane_index) + item_inputs = dict(broadcast_inputs) + for input_name, items in split_inputs.items(): + item_inputs[input_name] = items[item_index] + item_response = await invoke_remote_engine_async( + item_payload, + serialize_node_inputs(item_inputs), + ) + per_item_outputs[item_index] = deserialize_node_outputs(item_response) + completed_items += 1 + _emit_local_mapped_progress(payload, completed_items, total_items) + finally: + _clear_local_mapped_lane_progress(payload, lane_index, item_index) + + tasks = [asyncio.create_task(run_worker(lane_index)) for lane_index in range(parallelism)] + try: + await asyncio.gather(*tasks) + except Exception: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + return serialize_node_outputs( + _merge_static_and_mapped_outputs( + static_outputs=static_outputs, + mapped_outputs=_aggregate_mapped_outputs( + [item_outputs for item_outputs in per_item_outputs if item_outputs is not None], + { + **hybrid_payload, + "boundary_outputs": [ + boundary_output + for boundary_output in hybrid_payload.get("boundary_outputs", []) + if _is_mapped_boundary_output(boundary_output, hybrid_payload) + ], + }, + ), + payload=hybrid_payload, + ) ) - ) + finally: + if cleanup_payload is not None: + await invoke_remote_engine_async( + cleanup_payload, + serialize_node_inputs({}), + ) async def _invoke_mapped_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index beef9e7..fc6271a 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -3653,6 +3653,121 @@ def test_implicitly_mapped_subgraph_shared_model_keeps_unbatched_sampler_single_ ] +def test_implicitly_mapped_subgraph_clears_remote_session_once_after_all_items_finish( + remote_modal_app_module: Any, + serialization_module: Any, + session_state_module: Any, + monkeypatch: Any, +) -> None: + """Implicit mapped execution should reserve remote-session cleanup for one final cleanup payload.""" + observed_calls: list[tuple[str, bool, tuple[str, ...], dict[str, Any]]] = [] + session_cleared = False + + monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2) + + async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: + nonlocal session_cleared + hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload) + execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])) + clear_remote_session = bool(payload.get("clear_remote_session")) + observed_calls.append( + (str(payload["component_id"]), clear_remote_session, execute_node_ids, hydrated_inputs) + ) + + if clear_remote_session: + session_cleared = True + elif session_cleared: + raise session_state_module.RemoteSessionStateError( + "Remote session 'session-1' was not found." + ) + + if execute_node_ids == ("4",): + return serialization_module.serialize_node_outputs(("sampler-4",)) + if execute_node_ids == ("12",): + return serialization_module.serialize_node_outputs((f"sampler-12:{hydrated_inputs['remote_input_0']}",)) + if execute_node_ids == (): + return serialization_module.serialize_node_outputs(()) + raise AssertionError(f"Unexpected execute nodes for implicit cleanup regression: {execute_node_ids!r}") + + monkeypatch.setattr( + remote_modal_app_module, + "invoke_remote_engine_async", + fake_invoke_remote_engine_async, + ) + + payload = { + "payload_kind": "subgraph", + "component_id": "17", + "prompt_id": "prompt-1", + "component_node_ids": ["4", "12", "17"], + "execute_node_ids": ["4", "12"], + "subgraph_prompt": { + "17": {"class_type": "LoraLoaderModelOnly", "inputs": {}}, + "4": {"class_type": "KSampler", "inputs": {"model": ["17", 0], "seed": 0}}, + "12": { + "class_type": "KSampler", + "inputs": {"model": ["17", 0], "seed": 0}, + }, + }, + "boundary_inputs": [ + { + "proxy_input_name": "remote_input_0", + "io_type": "INT", + "targets": [{"node_id": "12", "input_name": "seed"}], + }, + { + "proxy_input_name": "remote_input_1", + "io_type": "INT", + "targets": [{"node_id": "4", "input_name": "seed"}], + }, + ], + "boundary_outputs": [ + {"node_id": "4", "io_type": "STRING", "is_list": False}, + {"node_id": "12", "io_type": "STRING", "is_list": False}, + ], + "extra_data": {"client_id": "client-1"}, + "remote_session": session_state_module.RemoteSessionHandle( + session_id="session-1", + prompt_id="prompt-1", + owner_component_id="17", + ).to_payload(), + "clear_remote_session": True, + } + + response = asyncio.run( + remote_modal_app_module._invoke_implicitly_mapped_subgraph_async( + payload, + serialization_module.serialize_node_inputs( + { + "remote_input_0": [10, 11, 12], + "remote_input_1": [28], + } + ), + ) + ) + + assert serialization_module.deserialize_node_outputs(response) == ( + "sampler-4", + ["sampler-12:10", "sampler-12:11", "sampler-12:12"], + ) + assert [ + (component_id, execute_node_ids) + for component_id, _, execute_node_ids, _ in observed_calls + ] == [ + ("17::static", ("4",)), + ("17::item:0", ("12",)), + ("17::item:1", ("12",)), + ("17::item:2", ("12",)), + ("17::cleanup", ()), + ] + assert all( + not clear_remote_session + for component_id, clear_remote_session, _, _ in observed_calls + if component_id != "17::cleanup" + ) + assert observed_calls[-1] == ("17::cleanup", True, (), {}) + + def test_implicitly_mapped_subgraph_keeps_conditioning_lists_broadcast( remote_modal_app_module: Any, serialization_module: Any,