diff --git a/README.md b/README.md index 40d839f..e1ea80c 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ Remote runtime behavior: - Hybrid mapped sub-runs now trim their prompt down to the dependency closure of the specific `execute_node_ids` they are about to run. That keeps a static-only sub-run from validating or executing the unrelated mapped branch just because both branches originally lived in the same coarse remote component. - Hybrid sub-run trimming also tolerates stale execute target ids that are no longer present in the current subgraph prompt. Those ids are now dropped before dependency resolution instead of crashing the remote worker with a `KeyError`. - Remote cancellation now uses a shared Modal `Dict` control store instead of a second RPC lane into the execution class, so per-container execution concurrency can stay at `1` without losing interrupt propagation. +- If a remote execution crashes instead of being interrupted normally, the worker now schedules its own container process to exit after surfacing the exception. That retires the bad warm GPU container instead of leaving it idle and billable after a hard failure. - If a run is cancelled and restarted quickly, or a warm container is still releasing heavy model files, the remote worker now gives Modal volume reload a longer bounded retry window so recently released files can close before the next request needs a fresh `vol.reload()`. - Mapped remote execution now carries a per-request volume reload marker, so one warm container only performs `vol.reload()` once for that uploaded asset snapshot even if the local scheduler fans the component out into many per-item Modal calls. - Volume reload dedupe is now request-wide, not component-wide. If one queued workflow rewrites into several remote components, they all share the same reload marker so a warm container only reloads once for that overall synced asset set instead of once per component. @@ -317,6 +318,7 @@ If a remote-marked node depends on a model filename that cannot be resolved to a - `COMFY_MODAL_EXECUTION_MODE`: Set to `local` for in-process fallback execution. Default: `local`. - `COMFY_MODAL_APP_NAME` and `COMFY_MODAL_VOLUME_NAME`: Override Modal app and volume naming. - `COMFY_MODAL_INTERRUPT_DICT_NAME`: Override the shared Modal `Dict` used for remote cancellation flags. Default: `-interrupts`. +- `COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR`: When true, a remote execution crash makes the worker exit its own container after returning the error. Default: `true`. - `COMFY_MODAL_AUTO_DEPLOY`: Automatically deploy the Modal app on first remote invocation when lookup fails. Default: `true`. - `COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK`: Re-enable slow `app.run()` fallback in remote mode. Default: `false`. diff --git a/api_intercept.py b/api_intercept.py index 7db53b0..74d4d31 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -1239,6 +1239,7 @@ def _build_component_payload( requires_volume_reload: bool, volume_reload_marker: str | None, custom_nodes_bundle: SyncedAsset | None, + terminate_container_on_error: bool, ) -> dict[str, Any]: """Build the serialized execution payload for one remote component.""" payload = { @@ -1286,6 +1287,7 @@ def _build_component_payload( "extra_data": copy.deepcopy(extra_data or {}), "requires_volume_reload": requires_volume_reload, "volume_reload_marker": volume_reload_marker, + "terminate_container_on_error": terminate_container_on_error, "custom_nodes_bundle": ( custom_nodes_bundle.remote_path if custom_nodes_bundle is not None else None ), @@ -1475,6 +1477,7 @@ def rewrite_prompt_for_modal( requires_volume_reload=requires_volume_reload, volume_reload_marker=volume_reload_marker, custom_nodes_bundle=summary.custom_nodes_bundle, + terminate_container_on_error=resolved_settings.terminate_container_on_error, ) _rewrite_component_into_proxy( component=component, diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 790f5a3..ac4954f 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -55,6 +55,7 @@ _LOADER_CACHE_WRAPPED_CLASSES: set[str] = set() _LOADER_OUTPUT_CACHE: dict[tuple[str, str], tuple[Any, ...]] = {} _PROMPT_EXECUTOR_STATES_LOCK = threading.Lock() _MODAL_VOLUME_RELOAD_MARKERS_LOCK = threading.Lock() +_CONTAINER_TERMINATION_LOCK = threading.Lock() try: import modal # type: ignore @@ -79,6 +80,8 @@ _MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS = (0.0, 0.25, 0.5, 1.0, 2.0, _MODAL_VOLUME_RELOAD_MARKER_CACHE_LIMIT = 256 _MODAL_VOLUME_RELOAD_MARKERS: queue.SimpleQueue[str] | None = None _MODAL_VOLUME_RELOAD_MARKER_SET: set[str] = set() +_REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS = 1.0 +_CONTAINER_TERMINATION_SCHEDULED = False @dataclass @@ -89,6 +92,57 @@ class _RemoteExecutionControl: interrupt_flag_key: str +def _schedule_process_exit(delay_seconds: float, exit_code: int) -> None: + """Exit the current process after a short delay to retire a bad Modal container.""" + + def exit_later() -> None: + """Sleep briefly so Modal can ship the error response before exiting the worker.""" + if delay_seconds > 0: + time.sleep(delay_seconds) + logger.error("Exiting Modal container process with code=%s after remote failure.", exit_code) + sys.stdout.flush() + sys.stderr.flush() + os._exit(exit_code) + + threading.Thread( + target=exit_later, + name="modal-container-exit", + daemon=True, + ).start() + + +def _is_interrupt_like_failure(exc: Exception) -> bool: + """Return whether one remote failure represents an expected interruption rather than a crash.""" + return "interrupt" in str(exc).lower() + + +def _maybe_schedule_container_termination_on_error( + payload: dict[str, Any], + exc: Exception, +) -> bool: + """Retire the current Modal container after a remote execution crash when configured.""" + if not _is_modal_container_runtime(): + return False + if not bool(payload.get("terminate_container_on_error", True)): + return False + if _is_interrupt_like_failure(exc): + return False + + global _CONTAINER_TERMINATION_SCHEDULED + with _CONTAINER_TERMINATION_LOCK: + if _CONTAINER_TERMINATION_SCHEDULED: + return False + _CONTAINER_TERMINATION_SCHEDULED = True + + logger.error( + "Scheduling Modal container termination after remote execution failure for component=%s.", + payload.get("component_id"), + exc_info=(type(exc), exc, exc.__traceback__), + ) + _schedule_process_exit(_REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS, 1) + return True + + class _NullPromptServer: """Minimal PromptExecutor server stub for headless subgraph execution.""" @@ -1898,35 +1952,39 @@ if modal is not None: # pragma: no branch - remote entrypoint configuration. """Execute a proxied node or subgraph inside the Modal container.""" component_id = payload.get("component_id", "single-node") reload_marker = _modal_volume_reload_marker(payload) - with _registered_remote_execution(payload) as execution_control: - with _timed_phase( - "remote_engine_execute_payload", - component=component_id, - payload_kind=payload.get("payload_kind"), - ): - if _should_reload_modal_volume(payload): - _reload_modal_volume_for_request( - vol, - str(component_id), - reload_marker=reload_marker, - ) - else: - _emit_modal_volume_reload_skip(component_id, payload) - if payload.get("payload_kind") == "subgraph": - return execute_subgraph_locally( + try: + with _registered_remote_execution(payload) as execution_control: + with _timed_phase( + "remote_engine_execute_payload", + component=component_id, + payload_kind=payload.get("payload_kind"), + ): + if _should_reload_modal_volume(payload): + _reload_modal_volume_for_request( + vol, + str(component_id), + reload_marker=reload_marker, + ) + else: + _emit_modal_volume_reload_skip(component_id, payload) + if payload.get("payload_kind") == "subgraph": + return execute_subgraph_locally( + payload, + kwargs_payload, + cancellation_event=execution_control.cancellation_event, + interrupt_store=interrupt_flags, + interrupt_flag_key=execution_control.interrupt_flag_key, + ) + return execute_node_locally( payload, kwargs_payload, cancellation_event=execution_control.cancellation_event, interrupt_store=interrupt_flags, interrupt_flag_key=execution_control.interrupt_flag_key, ) - return execute_node_locally( - payload, - kwargs_payload, - cancellation_event=execution_control.cancellation_event, - interrupt_store=interrupt_flags, - interrupt_flag_key=execution_control.interrupt_flag_key, - ) + except Exception as exc: + _maybe_schedule_container_termination_on_error(payload, exc) + raise @modal.method() def execute_payload_stream( @@ -1937,27 +1995,31 @@ if modal is not None: # pragma: no branch - remote entrypoint configuration. """Stream progress envelopes and a final serialized result for one payload.""" component_id = payload.get("component_id", "single-node") reload_marker = _modal_volume_reload_marker(payload) - with _registered_remote_execution(payload) as execution_control: - with _timed_phase( - "remote_engine_execute_payload", - component=component_id, - payload_kind=payload.get("payload_kind"), - ): - if _should_reload_modal_volume(payload): - _reload_modal_volume_for_request( - vol, - str(component_id), - reload_marker=reload_marker, + try: + with _registered_remote_execution(payload) as execution_control: + with _timed_phase( + "remote_engine_execute_payload", + component=component_id, + payload_kind=payload.get("payload_kind"), + ): + if _should_reload_modal_volume(payload): + _reload_modal_volume_for_request( + vol, + str(component_id), + reload_marker=reload_marker, + ) + else: + _emit_modal_volume_reload_skip(component_id, payload) + yield from _stream_remote_payload_events( + payload, + kwargs_payload, + cancellation_event=execution_control.cancellation_event, + interrupt_store=interrupt_flags, + interrupt_flag_key=execution_control.interrupt_flag_key, ) - else: - _emit_modal_volume_reload_skip(component_id, payload) - yield from _stream_remote_payload_events( - payload, - kwargs_payload, - cancellation_event=execution_control.cancellation_event, - interrupt_store=interrupt_flags, - interrupt_flag_key=execution_control.interrupt_flag_key, - ) + except Exception as exc: + _maybe_schedule_container_termination_on_error(payload, exc) + raise else: app = None diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 8ba9431..a6fda09 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -86,6 +86,7 @@ - Once the easy wrapper shapes are normalized, the next debugging step should be structured worker-side diagnostics, not more blind guesses. Logging the normalized execute-node ids, boundary outputs, suspicious remaining wrapped prompt inputs, and the raw ComfyUI `execution_error` payload makes the surviving malformed field visible on the very next remote failure. - Hybrid mapped components cannot hand the whole coarse component prompt to every sub-run. Even when a static sibling should execute once, `PromptExecutor` still validates the full prompt, so static and per-item sub-runs must prune the rewritten prompt to the dependency closure of their own `execute_node_ids` or the unrelated branch's batched inputs can still fail validation. - That trim pass also has to tolerate stale execute ids that survive into a sub-run payload after rewrite. Dependency resolution should drop any execute target missing from the current trimmed prompt instead of assuming every listed node id still exists and crashing with `KeyError`. +- Modal's public Python SDK does not expose a direct "kill this `@app.cls` container" call from the local client. If we want crashy warm GPU workers to disappear immediately, the practical design is to have the local side opt into crash teardown and the remote worker schedule `os._exit(1)` on itself after surfacing the exception. - Content-addressed sync should trust the deterministic payload path, not only the sidecar marker. If a crash happens after uploading a hash-named asset but before writing its `.done` marker, the next run should backfill the marker and reuse the existing payload instead of trying the same upload again and failing with `FileExistsError`. - The streamed Modal generator path needs the same tensor-aware serialization rules as ordinary node input/output transport. If a progress or result envelope leaks a raw `torch.Tensor`, Modal's JSON encoding fails before the local client ever sees it, so the stream layer should normalize event payloads through the shared serializer too. - Modal volume metadata calls can become a bottleneck on their own. Repeated `VolumeListFiles` checks for the same marker path are worth memoizing for the life of the queue request, and transient `ResourceExhaustedError` responses should be retried with backoff instead of surfacing immediately as hard prompt failures. diff --git a/settings.py b/settings.py index 12720d7..f1606b6 100644 --- a/settings.py +++ b/settings.py @@ -31,6 +31,7 @@ class ModalSyncSettings: comfyui_root: Path | None custom_nodes_dir: Path | None interrupt_dict_name: str = "comfy-modal-sync-interrupts" + terminate_container_on_error: bool = True modal_gpu: str = "A100" scaledown_window_seconds: int = 600 min_containers: int = 0 @@ -169,6 +170,8 @@ def get_settings() -> ModalSyncSettings: "COMFY_MODAL_INTERRUPT_DICT_NAME", f"{app_name}-interrupts", ), + terminate_container_on_error=_read_bool_env("COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR") + is not False, modal_gpu=os.getenv("COMFY_MODAL_GPU", "A100").strip() or "A100", scaledown_window_seconds=_read_int_env("COMFY_MODAL_SCALEDOWN_WINDOW", 600), min_containers=_read_int_env("COMFY_MODAL_MIN_CONTAINERS", 0), diff --git a/tests/conftest.py b/tests/conftest.py index 7763cf3..8e6785a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -133,6 +133,7 @@ def reset_modal_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> monkeypatch.delenv("COMFY_MODAL_ENABLE_MEMORY_SNAPSHOT", raising=False) monkeypatch.delenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", raising=False) monkeypatch.delenv("COMFY_MODAL_INTERRUPT_DICT_NAME", raising=False) + monkeypatch.delenv("COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR", raising=False) monkeypatch.delenv("COMFY_MODAL_MAX_CONTAINERS", raising=False) monkeypatch.delenv("COMFY_MODAL_BUFFER_CONTAINERS", raising=False) monkeypatch.delenv("COMFY_MODAL_SCALEDOWN_WINDOW", raising=False) diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 89df3f0..5789da2 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -208,6 +208,7 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy( assert payload["requires_volume_reload"] is True assert isinstance(payload["volume_reload_marker"], str) assert payload["volume_reload_marker"] + assert payload["terminate_container_on_error"] is True assert payload["boundary_inputs"] == [] assert payload["boundary_outputs"] == [ { diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index e6d74de..b5fcd6a 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -501,6 +501,60 @@ def test_modal_cloud_only_reloads_volume_for_requests_with_new_uploads( assert modal_cloud_module._should_reload_modal_volume({}) is True +def test_modal_cloud_schedules_container_exit_for_remote_failures( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """Unhandled remote execution failures should retire the current Modal container.""" + scheduled_exits: list[tuple[float, int]] = [] + original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED + monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False) + monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True) + monkeypatch.setattr( + modal_cloud_module, + "_schedule_process_exit", + lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)), + ) + try: + scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error( + {"component_id": "component-1", "terminate_container_on_error": True}, + RuntimeError("boom"), + ) + finally: + monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag) + + assert scheduled is True + assert scheduled_exits == [(modal_cloud_module._REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS, 1)] + + +def test_modal_cloud_does_not_schedule_container_exit_for_interruptions( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """Expected interruption-style failures should not tear down the warm Modal container.""" + scheduled_exits: list[tuple[float, int]] = [] + original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED + monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False) + monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True) + monkeypatch.setattr( + modal_cloud_module, + "_schedule_process_exit", + lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)), + ) + try: + scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error( + {"component_id": "component-1", "terminate_container_on_error": True}, + modal_cloud_module.RemoteSubgraphExecutionError( + "Remote subgraph execution was interrupted." + ), + ) + finally: + monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag) + + assert scheduled is False + assert scheduled_exits == [] + + def test_modal_cloud_skips_duplicate_reload_markers_in_same_container( modal_cloud_module: Any, ) -> None: diff --git a/tests/test_settings.py b/tests/test_settings.py index 463210a..1a9c88b 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -106,3 +106,18 @@ def test_settings_reads_interrupt_dict_name_override( settings_module.get_settings.cache_clear() assert settings.interrupt_dict_name == "custom-interrupt-store" + + +def test_settings_reads_terminate_container_on_error_override( + settings_module: Any, + monkeypatch: Any, +) -> None: + """Remote crash teardown should be configurable via environment variable.""" + monkeypatch.setenv("COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR", "false") + settings_module.get_settings.cache_clear() + try: + settings = settings_module.get_settings() + finally: + settings_module.get_settings.cache_clear() + + assert settings.terminate_container_on_error is False