1580 lines
60 KiB
Python
1580 lines
60 KiB
Python
"""Stable Modal cloud entrypoint for ComfyUI Modal-Sync."""
|
|
|
|
import asyncio
|
|
import copy
|
|
import gc
|
|
import hashlib
|
|
import importlib
|
|
import importlib.metadata
|
|
import importlib.util
|
|
from io import BytesIO
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import queue
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parent
|
|
_REMOTE_REPO_ROOT = Path("/root/comfyui_modal_sync_repo")
|
|
_LOCAL_COMFYUI_ROOT = (Path.home() / "git" / "ComfyUI").resolve()
|
|
_REMOTE_COMFYUI_ROOT = Path("/root/comfyui_src")
|
|
_REMOTE_LLM_COMPILE_CACHE_ROOT = Path("/root/.cache/comfy-modal-llm")
|
|
for candidate in (
|
|
_REPO_ROOT,
|
|
_REMOTE_REPO_ROOT,
|
|
_LOCAL_COMFYUI_ROOT,
|
|
_REMOTE_COMFYUI_ROOT,
|
|
):
|
|
candidate_str = str(candidate)
|
|
try:
|
|
candidate_exists = candidate.exists()
|
|
except PermissionError:
|
|
candidate_exists = False
|
|
if candidate_exists and candidate_str not in sys.path:
|
|
sys.path.insert(0, candidate_str)
|
|
|
|
from runtime_environment import ( # noqa: E402 - paths are bootstrapped above.
|
|
COMFYUI_RUNTIME_SOURCE_DIRECTORIES as _COMFYUI_IMAGE_RUNTIME_DIRECTORIES,
|
|
COMFYUI_RUNTIME_SOURCE_FILES as _COMFYUI_IMAGE_RUNTIME_FILES,
|
|
REMOTE_APP_PROTOCOL_VERSION as _REMOTE_APP_PROTOCOL_VERSION,
|
|
REMOTE_PYTHON_VERSION,
|
|
RemoteTorchBuild as _RemoteTorchBuild,
|
|
build_remote_runtime_identity,
|
|
custom_node_runtime_packages as _custom_node_runtime_packages,
|
|
remote_accelerator_packages as _remote_accelerator_packages,
|
|
remote_accelerator_validation_command as _remote_accelerator_validation_command,
|
|
remote_apt_packages as _comfyui_apt_packages,
|
|
remote_huggingface_packages as _remote_huggingface_packages,
|
|
remote_huggingface_validation_command as _remote_huggingface_validation_command,
|
|
remote_runtime_packages as _comfyui_runtime_packages,
|
|
select_remote_torch_build as _select_remote_torch_build,
|
|
)
|
|
from llm_recovery import ( # noqa: E402
|
|
LLM_FORCE_VLLM_THROUGHPUT_PAYLOAD_KEY,
|
|
is_llm_memory_recovery_exhausted,
|
|
)
|
|
from llm_staging import resolve_and_stage_model_references # noqa: E402
|
|
from durable_state import ( # noqa: E402 - paths are bootstrapped above.
|
|
DurableObjectCommitBatch,
|
|
DurableObjectRef,
|
|
DurableStateError,
|
|
FileDurableObjectStore,
|
|
InMemoryRemoteInvocationStore,
|
|
RemoteInvocationRecord,
|
|
new_running_invocation_record,
|
|
read_modal_volume_file,
|
|
)
|
|
from output_artifacts import ( # noqa: E402 - paths are bootstrapped above.
|
|
RemoteOutputSnapshot,
|
|
capture_execution_result,
|
|
snapshot_output_directory,
|
|
)
|
|
from remote_protocol import ( # noqa: E402 - paths are bootstrapped above.
|
|
BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
|
|
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
|
)
|
|
|
|
from serialization import ( # noqa: E402 - paths are bootstrapped above.
|
|
coerce_serialized_node_outputs,
|
|
deserialize_node_inputs,
|
|
deserialize_node_outputs,
|
|
deserialize_value,
|
|
serialize_mapping,
|
|
serialize_node_inputs,
|
|
serialize_node_outputs,
|
|
serialize_value,
|
|
)
|
|
from session_state import ( # noqa: E402 - paths are bootstrapped above.
|
|
InMemoryRemoteSessionBridgeStore,
|
|
InMemoryRemoteSessionStore,
|
|
RemoteSessionBridgeRecord,
|
|
RemoteSessionBridgeRecoveryKind,
|
|
RemoteSessionBridgeRef,
|
|
RemoteSessionHandle,
|
|
RemoteSessionStateError,
|
|
RemoteSessionValueRef,
|
|
is_remote_session_bridge_ref_payload,
|
|
is_remote_session_handle_payload,
|
|
is_remote_session_value_ref_payload,
|
|
stable_session_bridge_key,
|
|
)
|
|
from settings import ( # noqa: E402 - paths are bootstrapped above.
|
|
DEFAULT_MODAL_SECRET_NAME,
|
|
get_settings,
|
|
modal_deployment_app_name,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_comfy_bootstrap import (
|
|
CloudComfyBootstrapHooks,
|
|
_active_comfyui_root,
|
|
_alias_flux_rms_norm_weight_keys,
|
|
_build_checkpoint_loader_cache_key,
|
|
_build_clip_loader_cache_key,
|
|
_build_dual_clip_loader_cache_key,
|
|
_build_unet_loader_cache_key,
|
|
_build_vae_loader_cache_key,
|
|
_clone_loader_cache_outputs,
|
|
_clone_loader_cache_value,
|
|
_custom_node_package_for_candidate_file,
|
|
_ensure_comfy_runtime_initialized,
|
|
_ensure_comfyui_support_packages,
|
|
_ensure_default_custom_nodes_dir,
|
|
_ensure_headless_prompt_server_instance,
|
|
_ensure_prompt_node_classes_registered,
|
|
_extract_custom_nodes_bundle,
|
|
_force_import_package_from_root,
|
|
_install_loader_cache_wrappers,
|
|
_install_model_state_dict_compatibility_wrappers,
|
|
_iter_custom_nodes_manifest_assets,
|
|
_iter_missing_class_candidate_files,
|
|
_load_custom_nodes_manifest,
|
|
_load_execution_module,
|
|
_load_nodes_module,
|
|
_loader_cache_metric_snapshot,
|
|
_materialize_custom_nodes_manifest_assets,
|
|
_materialize_remote_asset_path,
|
|
_missing_node_class_diagnostics,
|
|
_patched_folder_paths_absolute_lookup,
|
|
_prompt_missing_node_class_types,
|
|
_readthrough_cache_path,
|
|
_record_loader_cache_metric,
|
|
_register_custom_nodes_root,
|
|
_register_modal_sync_runtime_nodes,
|
|
_reload_external_custom_nodes_for_missing_classes,
|
|
_resolve_custom_nodes_archives,
|
|
_resolve_custom_nodes_bundle_path,
|
|
_resolve_runtime_asset_path,
|
|
_rewrite_modal_asset_references,
|
|
_serialize_loader_cache_key,
|
|
_validated_custom_node_asset_relative_path,
|
|
_wrap_loader_method_with_cache,
|
|
configure_cloud_comfy_bootstrap_hooks,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_comfy_bootstrap import (
|
|
CloudComfyBootstrapHooks,
|
|
_active_comfyui_root,
|
|
_alias_flux_rms_norm_weight_keys,
|
|
_build_checkpoint_loader_cache_key,
|
|
_build_clip_loader_cache_key,
|
|
_build_dual_clip_loader_cache_key,
|
|
_build_unet_loader_cache_key,
|
|
_build_vae_loader_cache_key,
|
|
_clone_loader_cache_outputs,
|
|
_clone_loader_cache_value,
|
|
_custom_node_package_for_candidate_file,
|
|
_ensure_comfy_runtime_initialized,
|
|
_ensure_comfyui_support_packages,
|
|
_ensure_default_custom_nodes_dir,
|
|
_ensure_headless_prompt_server_instance,
|
|
_ensure_prompt_node_classes_registered,
|
|
_extract_custom_nodes_bundle,
|
|
_force_import_package_from_root,
|
|
_install_loader_cache_wrappers,
|
|
_install_model_state_dict_compatibility_wrappers,
|
|
_iter_custom_nodes_manifest_assets,
|
|
_iter_missing_class_candidate_files,
|
|
_load_custom_nodes_manifest,
|
|
_load_execution_module,
|
|
_load_nodes_module,
|
|
_loader_cache_metric_snapshot,
|
|
_materialize_custom_nodes_manifest_assets,
|
|
_materialize_remote_asset_path,
|
|
_missing_node_class_diagnostics,
|
|
_patched_folder_paths_absolute_lookup,
|
|
_prompt_missing_node_class_types,
|
|
_readthrough_cache_path,
|
|
_record_loader_cache_metric,
|
|
_register_custom_nodes_root,
|
|
_register_modal_sync_runtime_nodes,
|
|
_reload_external_custom_nodes_for_missing_classes,
|
|
_resolve_custom_nodes_archives,
|
|
_resolve_custom_nodes_bundle_path,
|
|
_resolve_runtime_asset_path,
|
|
_rewrite_modal_asset_references,
|
|
_serialize_loader_cache_key,
|
|
_validated_custom_node_asset_relative_path,
|
|
_wrap_loader_method_with_cache,
|
|
configure_cloud_comfy_bootstrap_hooks,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_node_output_cache import (
|
|
CloudNodeOutputCacheHooks,
|
|
_NodeOutputCacheLookupResult,
|
|
_PersistedNodeCacheRestoreState,
|
|
_await_maybe,
|
|
_boundary_output_node_ids,
|
|
_build_node_output_cache_immediate_signature,
|
|
_build_node_output_cache_signature_from_key_set_async,
|
|
_build_node_output_cache_signature_from_key_set_sync,
|
|
_cache_signature_link_output_index,
|
|
_canonicalize_node_output_cache_key_part,
|
|
_deserialize_node_output_cache_entry,
|
|
_emit_restored_node_cache_events,
|
|
_estimate_node_output_cache_value_size_bytes,
|
|
_include_unique_id_in_input_signature,
|
|
_install_prompt_executor_persisted_cache_restore,
|
|
_is_input_signature_cache_key_set,
|
|
_node_output_cache_ancestor_ids,
|
|
_node_output_cache_key,
|
|
_node_output_cache_key_from_key_set_async,
|
|
_node_output_cache_key_from_key_set_sync,
|
|
_node_output_cache_key_preview,
|
|
_node_output_cache_store,
|
|
_node_output_cache_store_get,
|
|
_node_output_cache_store_put,
|
|
_node_output_cache_value_preview,
|
|
_persist_node_output_cache_entries,
|
|
_prompt_executor_cache_get_sync,
|
|
_restore_persisted_node_output_cache_entries,
|
|
_restore_persisted_node_output_cache_entries_into_prepared_cache,
|
|
_serialize_node_output_cache_entry,
|
|
_tensor_cache_key_digest,
|
|
configure_cloud_node_output_cache_hooks,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_node_output_cache import (
|
|
CloudNodeOutputCacheHooks,
|
|
_NodeOutputCacheLookupResult,
|
|
_PersistedNodeCacheRestoreState,
|
|
_await_maybe,
|
|
_boundary_output_node_ids,
|
|
_build_node_output_cache_immediate_signature,
|
|
_build_node_output_cache_signature_from_key_set_async,
|
|
_build_node_output_cache_signature_from_key_set_sync,
|
|
_cache_signature_link_output_index,
|
|
_canonicalize_node_output_cache_key_part,
|
|
_deserialize_node_output_cache_entry,
|
|
_emit_restored_node_cache_events,
|
|
_estimate_node_output_cache_value_size_bytes,
|
|
_include_unique_id_in_input_signature,
|
|
_install_prompt_executor_persisted_cache_restore,
|
|
_is_input_signature_cache_key_set,
|
|
_node_output_cache_ancestor_ids,
|
|
_node_output_cache_key,
|
|
_node_output_cache_key_from_key_set_async,
|
|
_node_output_cache_key_from_key_set_sync,
|
|
_node_output_cache_key_preview,
|
|
_node_output_cache_store,
|
|
_node_output_cache_store_get,
|
|
_node_output_cache_store_put,
|
|
_node_output_cache_value_preview,
|
|
_persist_node_output_cache_entries,
|
|
_prompt_executor_cache_get_sync,
|
|
_restore_persisted_node_output_cache_entries,
|
|
_restore_persisted_node_output_cache_entries_into_prepared_cache,
|
|
_serialize_node_output_cache_entry,
|
|
_tensor_cache_key_digest,
|
|
configure_cloud_node_output_cache_hooks,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_prompt_execution import (
|
|
CloudPromptExecutionHooks,
|
|
_ReusablePromptExecutorState,
|
|
_aggregate_mapped_phase_outputs,
|
|
_apply_boundary_inputs,
|
|
_boundary_input_cache_signature,
|
|
_build_phase_subgraph_payload,
|
|
_coerce_primitive_prompt_input_value,
|
|
_coerce_prompt_primitive_input_values,
|
|
_collapse_cache_slot,
|
|
_copy_json_safe_prompt_metadata,
|
|
_execute_mapped_subgraph_payload,
|
|
_execute_node_locally_raw,
|
|
_execute_prompt_executor_compat,
|
|
_execute_subgraph_prompt,
|
|
_extract_prompt_executor_error,
|
|
_extract_prompt_executor_error_payload,
|
|
_format_prompt_executor_error_payload,
|
|
_get_or_create_prompt_executor_state,
|
|
_install_metadata_safe_dynamic_prompt_wrapper,
|
|
_invoke_original_node,
|
|
_is_link,
|
|
_log_prompt_executor_failure_details,
|
|
_mapped_phase_definition,
|
|
_merge_static_and_mapped_outputs,
|
|
_merge_static_or_mapped_values,
|
|
_node_input_type_map,
|
|
_node_input_types,
|
|
_node_required_input_names,
|
|
_normalize_link_output_index,
|
|
_normalize_prompt_input_value,
|
|
_normalize_subgraph_payload,
|
|
_prompt_executor_cache_config,
|
|
_prompt_executor_ram_thresholds,
|
|
_remote_session_ref_cache_signature,
|
|
_reset_prompt_executor_request_state,
|
|
_resolve_required_subgraph_nodes,
|
|
_serialize_prompt_executor_cache_scope,
|
|
_shared_subgraph_payload_fields,
|
|
_short_circuit_restored_session_output_subgraph,
|
|
_split_phase_outputs,
|
|
_summarize_suspicious_prompt_inputs,
|
|
_temporary_node_mapping,
|
|
_temporary_progress_hook,
|
|
_temporary_prompt_metadata,
|
|
_temporary_remote_interrupt_monitor,
|
|
_trim_subgraph_payload_to_required_nodes,
|
|
_unwrap_wrapped_prompt_link,
|
|
_validate_prompt_input_shapes,
|
|
_validate_required_prompt_inputs,
|
|
configure_cloud_prompt_execution_hooks,
|
|
execute_node_locally,
|
|
execute_subgraph_locally,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_prompt_execution import (
|
|
CloudPromptExecutionHooks,
|
|
_ReusablePromptExecutorState,
|
|
_aggregate_mapped_phase_outputs,
|
|
_apply_boundary_inputs,
|
|
_boundary_input_cache_signature,
|
|
_build_phase_subgraph_payload,
|
|
_coerce_primitive_prompt_input_value,
|
|
_coerce_prompt_primitive_input_values,
|
|
_collapse_cache_slot,
|
|
_copy_json_safe_prompt_metadata,
|
|
_execute_mapped_subgraph_payload,
|
|
_execute_node_locally_raw,
|
|
_execute_prompt_executor_compat,
|
|
_execute_subgraph_prompt,
|
|
_extract_prompt_executor_error,
|
|
_extract_prompt_executor_error_payload,
|
|
_format_prompt_executor_error_payload,
|
|
_get_or_create_prompt_executor_state,
|
|
_install_metadata_safe_dynamic_prompt_wrapper,
|
|
_invoke_original_node,
|
|
_is_link,
|
|
_log_prompt_executor_failure_details,
|
|
_mapped_phase_definition,
|
|
_merge_static_and_mapped_outputs,
|
|
_merge_static_or_mapped_values,
|
|
_node_input_type_map,
|
|
_node_input_types,
|
|
_node_required_input_names,
|
|
_normalize_link_output_index,
|
|
_normalize_prompt_input_value,
|
|
_normalize_subgraph_payload,
|
|
_prompt_executor_cache_config,
|
|
_prompt_executor_ram_thresholds,
|
|
_remote_session_ref_cache_signature,
|
|
_reset_prompt_executor_request_state,
|
|
_resolve_required_subgraph_nodes,
|
|
_serialize_prompt_executor_cache_scope,
|
|
_shared_subgraph_payload_fields,
|
|
_short_circuit_restored_session_output_subgraph,
|
|
_split_phase_outputs,
|
|
_summarize_suspicious_prompt_inputs,
|
|
_temporary_node_mapping,
|
|
_temporary_progress_hook,
|
|
_temporary_prompt_metadata,
|
|
_temporary_remote_interrupt_monitor,
|
|
_trim_subgraph_payload_to_required_nodes,
|
|
_unwrap_wrapped_prompt_link,
|
|
_validate_prompt_input_shapes,
|
|
_validate_required_prompt_inputs,
|
|
configure_cloud_prompt_execution_hooks,
|
|
execute_node_locally,
|
|
execute_subgraph_locally,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_prompt_server_shims import (
|
|
CloudPromptServerHooks,
|
|
_HeadlessPromptQueue,
|
|
_HeadlessPromptServerInstance,
|
|
_NullPromptServer,
|
|
_TracingPromptServer,
|
|
configure_cloud_prompt_server_hooks,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_prompt_server_shims import (
|
|
CloudPromptServerHooks,
|
|
_HeadlessPromptQueue,
|
|
_HeadlessPromptServerInstance,
|
|
_NullPromptServer,
|
|
_TracingPromptServer,
|
|
configure_cloud_prompt_server_hooks,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_streaming import (
|
|
CloudStreamingErrors,
|
|
_BoundedStreamEventBuffer,
|
|
_abandon_streamed_remote_invocation,
|
|
_stream_remote_payload_events as _stream_remote_payload_events_impl,
|
|
configure_cloud_streaming_errors,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_streaming import (
|
|
CloudStreamingErrors,
|
|
_BoundedStreamEventBuffer,
|
|
_abandon_streamed_remote_invocation,
|
|
_stream_remote_payload_events as _stream_remote_payload_events_impl,
|
|
configure_cloud_streaming_errors,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_volume_reload import (
|
|
CloudVolumeReloadHooks,
|
|
_MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS,
|
|
_clear_warm_remote_caches,
|
|
_custom_nodes_manifest_dependency_paths,
|
|
_download_committed_volume_path,
|
|
_emit_modal_volume_reload_skip,
|
|
_has_seen_modal_volume_reload_marker,
|
|
_hydrate_missing_payload_volume_paths,
|
|
_is_modal_volume_open_files_error,
|
|
_iter_payload_input_strings,
|
|
_log_payload_volume_reload_diagnostics,
|
|
_modal_volume_reload_marker,
|
|
_payload_uploaded_volume_paths,
|
|
_payload_uploaded_volume_paths_visible,
|
|
_payload_volume_paths,
|
|
_payload_volume_paths_visible,
|
|
_prepare_for_modal_volume_reload,
|
|
_record_modal_volume_reload_marker,
|
|
_reload_modal_volume_for_request,
|
|
_runtime_volume_path_visible,
|
|
_should_reload_modal_volume,
|
|
_sleep_before_modal_volume_reload_retry,
|
|
configure_cloud_volume_reload_hooks,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_volume_reload import (
|
|
CloudVolumeReloadHooks,
|
|
_MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS,
|
|
_clear_warm_remote_caches,
|
|
_custom_nodes_manifest_dependency_paths,
|
|
_download_committed_volume_path,
|
|
_emit_modal_volume_reload_skip,
|
|
_has_seen_modal_volume_reload_marker,
|
|
_hydrate_missing_payload_volume_paths,
|
|
_is_modal_volume_open_files_error,
|
|
_iter_payload_input_strings,
|
|
_log_payload_volume_reload_diagnostics,
|
|
_modal_volume_reload_marker,
|
|
_payload_uploaded_volume_paths,
|
|
_payload_uploaded_volume_paths_visible,
|
|
_payload_volume_paths,
|
|
_payload_volume_paths_visible,
|
|
_prepare_for_modal_volume_reload,
|
|
_record_modal_volume_reload_marker,
|
|
_reload_modal_volume_for_request,
|
|
_runtime_volume_path_visible,
|
|
_should_reload_modal_volume,
|
|
_sleep_before_modal_volume_reload_retry,
|
|
configure_cloud_volume_reload_hooks,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_prewarm import (
|
|
CloudPrewarmHooks,
|
|
_LLMCompileMissCheckpoint,
|
|
_build_loader_prewarm_payload,
|
|
_commit_actual_llm_compile_cache,
|
|
_execute_llm_prewarm_plans,
|
|
_execute_loader_prewarm_plans,
|
|
_llm_compile_manifest_path,
|
|
_llm_compile_miss_checkpoint,
|
|
_llm_profiles_in_payload,
|
|
_llm_prewarm_model_profile,
|
|
_loader_prewarm_plan_key,
|
|
_log_compile_cache_memory_maps,
|
|
_mapped_process_files_under,
|
|
_prepare_warm_container_for_request,
|
|
_prewarm_restored_runtime,
|
|
_prewarm_snapshot_state,
|
|
_reload_compile_cache_volume,
|
|
_triton_compile_listener_engine_pids,
|
|
_triton_compile_miss_signal_size,
|
|
_write_llm_compile_manifest,
|
|
configure_cloud_prewarm_hooks,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_prewarm import (
|
|
CloudPrewarmHooks,
|
|
_LLMCompileMissCheckpoint,
|
|
_build_loader_prewarm_payload,
|
|
_commit_actual_llm_compile_cache,
|
|
_execute_llm_prewarm_plans,
|
|
_execute_loader_prewarm_plans,
|
|
_llm_compile_manifest_path,
|
|
_llm_compile_miss_checkpoint,
|
|
_llm_profiles_in_payload,
|
|
_llm_prewarm_model_profile,
|
|
_loader_prewarm_plan_key,
|
|
_log_compile_cache_memory_maps,
|
|
_mapped_process_files_under,
|
|
_prepare_warm_container_for_request,
|
|
_prewarm_restored_runtime,
|
|
_prewarm_snapshot_state,
|
|
_reload_compile_cache_volume,
|
|
_triton_compile_listener_engine_pids,
|
|
_triton_compile_miss_signal_size,
|
|
_write_llm_compile_manifest,
|
|
configure_cloud_prewarm_hooks,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_session_bridge import (
|
|
CloudSessionBridgeHooks,
|
|
_RemoteSessionBridgeResolutionStats,
|
|
_bridge_record_replays_sampling_node,
|
|
_build_durable_bridge_rehydration_plan,
|
|
_build_remote_session_bridge_record,
|
|
_deserialize_remote_session_bridge_producer_inputs,
|
|
_get_remote_session_bridge_value,
|
|
_json_payload_size_bytes,
|
|
_load_loader_snapshot_profile,
|
|
_load_remote_session_bridge_record,
|
|
_log_remote_session_resolution_summary,
|
|
_offload_large_bridge_payloads,
|
|
_payload_remote_session_handle,
|
|
_record_remote_session_resolution_event,
|
|
_rehydrate_remote_session_bridge_value,
|
|
_remote_session_bridge_recovery_input_names,
|
|
_remote_session_bridge_replay_stack,
|
|
_resolve_remote_session_inputs,
|
|
_restore_planned_remote_session_bridge_value,
|
|
_restore_serialized_remote_session_bridge_value,
|
|
_sanitize_payload_for_session_bridge_record,
|
|
_select_remote_session_bridge_recovery_kind,
|
|
_serialize_durable_bridge_output,
|
|
_session_bridge_store,
|
|
_snapshot_profile_store,
|
|
_store_remote_session_bridge_record,
|
|
_store_remote_session_bridge_value,
|
|
_subgraph_contains_sampling_node,
|
|
configure_cloud_session_bridge_hooks,
|
|
remote_session_store,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_session_bridge import (
|
|
CloudSessionBridgeHooks,
|
|
_RemoteSessionBridgeResolutionStats,
|
|
_bridge_record_replays_sampling_node,
|
|
_build_durable_bridge_rehydration_plan,
|
|
_build_remote_session_bridge_record,
|
|
_deserialize_remote_session_bridge_producer_inputs,
|
|
_get_remote_session_bridge_value,
|
|
_json_payload_size_bytes,
|
|
_load_loader_snapshot_profile,
|
|
_load_remote_session_bridge_record,
|
|
_log_remote_session_resolution_summary,
|
|
_offload_large_bridge_payloads,
|
|
_payload_remote_session_handle,
|
|
_record_remote_session_resolution_event,
|
|
_rehydrate_remote_session_bridge_value,
|
|
_remote_session_bridge_recovery_input_names,
|
|
_remote_session_bridge_replay_stack,
|
|
_resolve_remote_session_inputs,
|
|
_restore_planned_remote_session_bridge_value,
|
|
_restore_serialized_remote_session_bridge_value,
|
|
_sanitize_payload_for_session_bridge_record,
|
|
_select_remote_session_bridge_recovery_kind,
|
|
_serialize_durable_bridge_output,
|
|
_session_bridge_store,
|
|
_snapshot_profile_store,
|
|
_store_remote_session_bridge_record,
|
|
_store_remote_session_bridge_value,
|
|
_subgraph_contains_sampling_node,
|
|
configure_cloud_session_bridge_hooks,
|
|
remote_session_store,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_durable_invocation import (
|
|
DurableInvocationErrors,
|
|
_begin_remote_invocation,
|
|
_canary_barrier_marker_exists,
|
|
_canary_barrier_marker_key,
|
|
_canary_interrupt_requested,
|
|
_complete_remote_invocation,
|
|
_durable_object_store,
|
|
_execute_canary_payload,
|
|
_execute_payload_with_output_capture,
|
|
_execute_with_durable_invocation,
|
|
_fail_remote_invocation,
|
|
_invocation_record_store,
|
|
_load_completed_remote_invocation_result,
|
|
_load_remote_invocation_record,
|
|
_put_canary_barrier_marker,
|
|
_raise_if_canary_interrupted,
|
|
_remote_comfy_output_directory,
|
|
_remote_output_snapshot,
|
|
_store_remote_invocation_record,
|
|
_wait_for_canary_barrier,
|
|
_wait_for_running_remote_invocation,
|
|
configure_durable_invocation_errors,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_durable_invocation import (
|
|
DurableInvocationErrors,
|
|
_begin_remote_invocation,
|
|
_canary_barrier_marker_exists,
|
|
_canary_barrier_marker_key,
|
|
_canary_interrupt_requested,
|
|
_complete_remote_invocation,
|
|
_durable_object_store,
|
|
_execute_canary_payload,
|
|
_execute_payload_with_output_capture,
|
|
_execute_with_durable_invocation,
|
|
_fail_remote_invocation,
|
|
_invocation_record_store,
|
|
_load_completed_remote_invocation_result,
|
|
_load_remote_invocation_record,
|
|
_put_canary_barrier_marker,
|
|
_raise_if_canary_interrupted,
|
|
_remote_comfy_output_directory,
|
|
_remote_output_snapshot,
|
|
_store_remote_invocation_record,
|
|
_wait_for_canary_barrier,
|
|
_wait_for_running_remote_invocation,
|
|
configure_durable_invocation_errors,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_image_env import (
|
|
_REMOTE_LLM_COMPILE_CACHE_ROOT,
|
|
_install_custom_node_packages,
|
|
_install_remote_accelerator_packages,
|
|
_install_remote_torch_build,
|
|
_modal_image_environment,
|
|
_modal_secret_from_settings,
|
|
_model_stager_image_environment,
|
|
_remote_engine_cls_options,
|
|
_should_ignore_comfyui_path,
|
|
_should_ignore_repo_path,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_image_env import (
|
|
_REMOTE_LLM_COMPILE_CACHE_ROOT,
|
|
_install_custom_node_packages,
|
|
_install_remote_accelerator_packages,
|
|
_install_remote_torch_build,
|
|
_modal_image_environment,
|
|
_modal_secret_from_settings,
|
|
_model_stager_image_environment,
|
|
_remote_engine_cls_options,
|
|
_should_ignore_comfyui_path,
|
|
_should_ignore_repo_path,
|
|
)
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_app_guard import guard_against_existing_modal_app
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_app_guard import guard_against_existing_modal_app
|
|
try: # noqa: E402 - support package and flat Modal-container imports.
|
|
from .cloud_runtime_context import (
|
|
interrupt_flag_store,
|
|
invocation_record_store,
|
|
node_output_cache_store,
|
|
register_cloud_runtime_stores,
|
|
session_bridge_store,
|
|
snapshot_profile_store,
|
|
volume_store,
|
|
)
|
|
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
|
from cloud_runtime_context import (
|
|
interrupt_flag_store,
|
|
invocation_record_store,
|
|
node_output_cache_store,
|
|
register_cloud_runtime_stores,
|
|
session_bridge_store,
|
|
snapshot_profile_store,
|
|
volume_store,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Cloud app and ComfyUI bootstrap state.
|
|
_CLOUD_HANDLER_NAME = "comfyui-modal-sync-cloud-timestamped"
|
|
|
|
# Poisoned-container retirement state.
|
|
_CONTAINER_TERMINATION_LOCK = threading.Lock()
|
|
_REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS = 1.0
|
|
_CONTAINER_TERMINATION_SCHEDULED = False
|
|
try:
|
|
import modal # type: ignore
|
|
except ModuleNotFoundError: # pragma: no cover - remote entrypoint only.
|
|
modal = None
|
|
|
|
|
|
class RemoteSubgraphExecutionError(RuntimeError):
|
|
"""Raised when remote subgraph execution fails."""
|
|
|
|
|
|
class RemoteInvocationInProgressError(RuntimeError):
|
|
"""Raised when an idempotent invocation is already running remotely."""
|
|
|
|
|
|
class RemoteInvocationAbandonedError(RuntimeError):
|
|
"""Raised when a streamed invocation loses its consumer before completion."""
|
|
|
|
|
|
class RemoteCanaryInterruptedError(RuntimeError):
|
|
"""Raised when a live remote canary observes its shared interrupt flag."""
|
|
|
|
|
|
class RemoteCanaryBarrierTimeoutError(TimeoutError):
|
|
"""Raised when live canary calls fail to overlap before their deadline."""
|
|
|
|
|
|
class ExistingModalAppError(RuntimeError):
|
|
"""Raised when deploying would overwrite an existing Modal app."""
|
|
|
|
|
|
configure_durable_invocation_errors(
|
|
DurableInvocationErrors(
|
|
invocation_in_progress=RemoteInvocationInProgressError,
|
|
canary_interrupted=RemoteCanaryInterruptedError,
|
|
canary_barrier_timeout=RemoteCanaryBarrierTimeoutError,
|
|
)
|
|
)
|
|
|
|
configure_cloud_streaming_errors(
|
|
CloudStreamingErrors(invocation_abandoned=RemoteInvocationAbandonedError)
|
|
)
|
|
|
|
|
|
class RemoteFailureDisposition(str, Enum):
|
|
"""Describe whether one remote failure implies poisoned worker state."""
|
|
|
|
EXPECTED = "expected"
|
|
DETERMINISTIC = "deterministic"
|
|
POISONED_WORKER = "poisoned-worker"
|
|
|
|
|
|
def _guard_against_existing_modal_app(settings: Any, modal_module: Any) -> None:
|
|
"""Fail local Modal app construction when the configured app already exists."""
|
|
guard_against_existing_modal_app(
|
|
settings,
|
|
modal_module,
|
|
error_type=ExistingModalAppError,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _RemoteExecutionControl:
|
|
"""Track interruption state for one active remote payload execution."""
|
|
|
|
cancellation_event: threading.Event
|
|
interrupt_flag_key: str
|
|
|
|
|
|
def _meaningful_progress_values(
|
|
node_state: dict[str, Any]
|
|
) -> tuple[float, float] | None:
|
|
"""Return numeric progress values only for node states that represent real progress."""
|
|
try:
|
|
progress_value = float(node_state.get("value", 0.0))
|
|
max_value = float(node_state.get("max", 1.0))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
if max_value <= 1.0:
|
|
return None
|
|
return progress_value, max_value
|
|
|
|
|
|
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 _schedule_process_exit_unless_cancelled(
|
|
*,
|
|
delay_seconds: float,
|
|
exit_code: int,
|
|
cancel_event: threading.Event,
|
|
reason: str,
|
|
) -> None:
|
|
"""Exit the current process after a delay unless the caller cancels first."""
|
|
|
|
def exit_later() -> None:
|
|
"""Wait for cancellation or exit the worker if the delay expires."""
|
|
if delay_seconds > 0 and cancel_event.wait(timeout=delay_seconds):
|
|
logger.debug("Cancelled delayed Modal container restart for %s.", reason)
|
|
return
|
|
logger.error(
|
|
"Exiting Modal container process with code=%s after %s.", exit_code, reason
|
|
)
|
|
sys.stdout.flush()
|
|
sys.stderr.flush()
|
|
os._exit(exit_code)
|
|
|
|
threading.Thread(
|
|
target=exit_later,
|
|
name="modal-container-cancel-restart",
|
|
daemon=True,
|
|
).start()
|
|
|
|
|
|
def _schedule_remote_cancel_restart(
|
|
*,
|
|
component_id: str,
|
|
completion_event: threading.Event,
|
|
) -> bool:
|
|
"""Restart the Modal worker if a cancelled remote prompt keeps executing."""
|
|
if not _is_modal_container_runtime():
|
|
return False
|
|
|
|
delay_seconds = max(0.0, get_settings().remote_cancel_restart_seconds)
|
|
logger.warning(
|
|
"Remote cancellation requested for component=%s; scheduling container restart in %.3fs unless execution stops first.",
|
|
component_id,
|
|
delay_seconds,
|
|
)
|
|
_schedule_process_exit_unless_cancelled(
|
|
delay_seconds=delay_seconds,
|
|
exit_code=0,
|
|
cancel_event=completion_event,
|
|
reason=f"remote cancellation timeout for component={component_id}",
|
|
)
|
|
return True
|
|
|
|
|
|
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 _is_session_state_like_failure(exc: Exception) -> bool:
|
|
"""Return whether one remote failure came from prompt-scoped session routing/state issues."""
|
|
if isinstance(exc, RemoteSessionStateError):
|
|
return True
|
|
return "remote session" in str(exc).lower()
|
|
|
|
|
|
def _remote_failure_disposition(exc: Exception) -> RemoteFailureDisposition:
|
|
"""Classify one execution failure for worker-retirement decisions."""
|
|
if _is_interrupt_like_failure(exc) or _is_session_state_like_failure(exc):
|
|
return RemoteFailureDisposition.EXPECTED
|
|
if isinstance(exc, MemoryError) or is_llm_memory_recovery_exhausted(exc):
|
|
return RemoteFailureDisposition.POISONED_WORKER
|
|
|
|
message = str(exc).lower()
|
|
poisoned_runtime_markers = (
|
|
"cuda out of memory",
|
|
"cuda error",
|
|
"device-side assert",
|
|
"illegal memory access",
|
|
"cublas_status",
|
|
"cudnn_status",
|
|
"hip error",
|
|
)
|
|
if any(marker in message for marker in poisoned_runtime_markers):
|
|
return RemoteFailureDisposition.POISONED_WORKER
|
|
return RemoteFailureDisposition.DETERMINISTIC
|
|
|
|
|
|
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
|
|
disposition = _remote_failure_disposition(exc)
|
|
if disposition is not RemoteFailureDisposition.POISONED_WORKER:
|
|
logger.warning(
|
|
"Skipping Modal container termination for component=%s failure_disposition=%s because the worker is safe to reuse.",
|
|
payload.get("component_id"),
|
|
disposition.value,
|
|
exc_info=(type(exc), exc, exc.__traceback__),
|
|
)
|
|
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
|
|
|
|
|
|
|
|
|
|
def _build_cloud_log_formatter() -> logging.Formatter:
|
|
"""Return the default formatter for remote Modal-Sync logs with timestamps."""
|
|
return logging.Formatter(
|
|
fmt="%(asctime)s.%(msecs)03d +%(relativeCreated)07.0fms %(levelname)s [%(name)s] %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
)
|
|
|
|
|
|
def _configure_cloud_logging() -> logging.Logger:
|
|
"""Install a dedicated timestamped handler for the cloud runtime logger."""
|
|
logger.setLevel(logging.INFO)
|
|
for existing_handler in logger.handlers:
|
|
if getattr(existing_handler, "name", "") == _CLOUD_HANDLER_NAME:
|
|
return logger
|
|
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
handler.set_name(_CLOUD_HANDLER_NAME)
|
|
handler.setLevel(logging.INFO)
|
|
handler.setFormatter(_build_cloud_log_formatter())
|
|
logger.addHandler(handler)
|
|
logger.propagate = False
|
|
return logger
|
|
|
|
|
|
def _is_modal_container_runtime() -> bool:
|
|
"""Return whether the current process is executing inside a Modal container."""
|
|
return os.getenv("MODAL_IS_REMOTE") == "1" or bool(os.getenv("MODAL_TASK_ID"))
|
|
|
|
|
|
def _cloud_formatter() -> logging.Formatter:
|
|
"""Return the configured formatter used for cloud phase trace lines."""
|
|
for existing_handler in logger.handlers:
|
|
if getattr(existing_handler, "name", "") == _CLOUD_HANDLER_NAME:
|
|
formatter = existing_handler.formatter
|
|
if formatter is not None:
|
|
return formatter
|
|
return _build_cloud_log_formatter()
|
|
|
|
|
|
def _emit_cloud_info(message: str, *args: Any) -> None:
|
|
"""Emit an info line through logging and mirror it to stdout inside Modal containers."""
|
|
if not _is_modal_container_runtime():
|
|
logger.info(message, *args)
|
|
return
|
|
|
|
record = logger.makeRecord(
|
|
logger.name,
|
|
logging.INFO,
|
|
__file__,
|
|
0,
|
|
message,
|
|
args,
|
|
exc_info=None,
|
|
)
|
|
print(_cloud_formatter().format(record), file=sys.stdout, flush=True)
|
|
|
|
|
|
def _remote_execution_key(payload: dict[str, Any]) -> tuple[str, str]:
|
|
"""Return the registry key for one active remote execution."""
|
|
prompt_id = str(
|
|
payload.get("prompt_id") or payload.get("component_id") or "modal-subgraph"
|
|
)
|
|
component_id = str(payload.get("component_id") or "single-node")
|
|
return prompt_id, component_id
|
|
|
|
|
|
def _observe_remote_workflow_for_llm_mode(payload: dict[str, Any]) -> None:
|
|
"""Record real workflow arrivals for container-local vLLM auto promotion."""
|
|
if payload.get("payload_kind") == "canary":
|
|
return
|
|
from modal_llm_runtime import (
|
|
force_modal_vllm_throughput_after_memory_recovery,
|
|
observe_modal_workflow_execution,
|
|
)
|
|
|
|
prompt_id = payload.get("prompt_id")
|
|
normalized_prompt_id = str(prompt_id).strip() if prompt_id is not None else None
|
|
if bool(payload.get(LLM_FORCE_VLLM_THROUGHPUT_PAYLOAD_KEY)):
|
|
force_modal_vllm_throughput_after_memory_recovery(normalized_prompt_id)
|
|
return
|
|
observe_modal_workflow_execution(normalized_prompt_id)
|
|
|
|
|
|
def _remote_interrupt_flag_key(prompt_id: str, component_id: str) -> str:
|
|
"""Return the shared Modal interrupt-store key for one payload execution."""
|
|
return f"{prompt_id}:{component_id}"
|
|
|
|
|
|
@contextmanager
|
|
def _registered_remote_execution(
|
|
payload: dict[str, Any],
|
|
) -> Iterator[_RemoteExecutionControl]:
|
|
"""Prepare interruption state for one active remote execution."""
|
|
prompt_id, component_id = _remote_execution_key(payload)
|
|
control = _RemoteExecutionControl(
|
|
cancellation_event=threading.Event(),
|
|
interrupt_flag_key=_remote_interrupt_flag_key(prompt_id, component_id),
|
|
)
|
|
try:
|
|
yield control
|
|
finally:
|
|
interrupt_store = interrupt_flag_store()
|
|
if modal is not None and interrupt_store is not None:
|
|
interrupt_store.pop(control.interrupt_flag_key, None)
|
|
|
|
|
|
@contextmanager
|
|
def _timed_phase(phase: str, **fields: Any) -> Iterator[None]:
|
|
"""Log a start/finish pair with elapsed time for a named execution phase."""
|
|
field_suffix = ""
|
|
if fields:
|
|
rendered_fields = " ".join(f"{key}={value}" for key, value in fields.items())
|
|
field_suffix = f" {rendered_fields}"
|
|
phase_started_at = time.perf_counter()
|
|
_emit_cloud_info("Starting %s%s", phase, field_suffix)
|
|
try:
|
|
yield
|
|
finally:
|
|
_emit_cloud_info(
|
|
"Finished %s in %.3fs%s",
|
|
phase,
|
|
time.perf_counter() - phase_started_at,
|
|
field_suffix,
|
|
)
|
|
|
|
|
|
_configure_cloud_logging()
|
|
|
|
|
|
def _stream_remote_payload_events(
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes | bytearray | str | dict[str, Any],
|
|
cancellation_event: threading.Event | None = None,
|
|
interrupt_store: Any | None = None,
|
|
interrupt_flag_key: str | None = None,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Delegate remote event streaming through the stable cloud entrypoint name."""
|
|
yield from _stream_remote_payload_events_impl(
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event=cancellation_event,
|
|
interrupt_store=interrupt_store,
|
|
interrupt_flag_key=interrupt_flag_key,
|
|
)
|
|
|
|
|
|
configure_cloud_session_bridge_hooks(
|
|
CloudSessionBridgeHooks(
|
|
clone_loader_cache_value=_clone_loader_cache_value,
|
|
emit_cloud_info=_emit_cloud_info,
|
|
execute_node_locally_raw=_execute_node_locally_raw,
|
|
execute_subgraph_prompt=_execute_subgraph_prompt,
|
|
is_link=_is_link,
|
|
normalize_prompt_input_value=_normalize_prompt_input_value,
|
|
resolve_required_subgraph_nodes=_resolve_required_subgraph_nodes,
|
|
)
|
|
)
|
|
|
|
configure_cloud_comfy_bootstrap_hooks(
|
|
CloudComfyBootstrapHooks(
|
|
emit_cloud_info=_emit_cloud_info,
|
|
timed_phase=_timed_phase,
|
|
remote_subgraph_error=RemoteSubgraphExecutionError,
|
|
)
|
|
)
|
|
|
|
configure_cloud_node_output_cache_hooks(
|
|
CloudNodeOutputCacheHooks(
|
|
emit_cloud_info=_emit_cloud_info,
|
|
timed_phase=_timed_phase,
|
|
)
|
|
)
|
|
|
|
configure_cloud_volume_reload_hooks(
|
|
CloudVolumeReloadHooks(
|
|
emit_cloud_info=_emit_cloud_info,
|
|
timed_phase=_timed_phase,
|
|
)
|
|
)
|
|
|
|
configure_cloud_prewarm_hooks(
|
|
CloudPrewarmHooks(
|
|
emit_cloud_info=_emit_cloud_info,
|
|
timed_phase=_timed_phase,
|
|
)
|
|
)
|
|
|
|
configure_cloud_prompt_execution_hooks(
|
|
CloudPromptExecutionHooks(
|
|
emit_cloud_info=_emit_cloud_info,
|
|
timed_phase=_timed_phase,
|
|
schedule_remote_cancel_restart=_schedule_remote_cancel_restart,
|
|
remote_subgraph_error=RemoteSubgraphExecutionError,
|
|
)
|
|
)
|
|
|
|
configure_cloud_prompt_server_hooks(
|
|
CloudPromptServerHooks(
|
|
collapse_cache_slot=_collapse_cache_slot,
|
|
emit_cloud_info=_emit_cloud_info,
|
|
meaningful_progress_values=_meaningful_progress_values,
|
|
)
|
|
)
|
|
|
|
|
|
if modal is not None: # pragma: no branch - remote entrypoint configuration.
|
|
settings = globals().get("__comfy_modal_settings_override__") or get_settings()
|
|
__comfy_modal_gpu__ = settings.modal_gpu
|
|
__comfy_modal_app_name__ = modal_deployment_app_name(settings)
|
|
__comfy_modal_secret_name__ = str(
|
|
getattr(settings, "modal_secret_name", DEFAULT_MODAL_SECRET_NAME)
|
|
).strip()
|
|
_guard_against_existing_modal_app(settings, modal)
|
|
app = modal.App(__comfy_modal_app_name__)
|
|
modal_secret = _modal_secret_from_settings(settings, modal)
|
|
vol = modal.Volume.from_name(settings.volume_name, create_if_missing=True)
|
|
llm_compile_cache_vol = modal.Volume.from_name(
|
|
getattr(
|
|
settings,
|
|
"llm_compile_cache_volume_name",
|
|
f"{settings.volume_name}-llm-compile-cache",
|
|
),
|
|
create_if_missing=True,
|
|
)
|
|
interrupt_flags = modal.Dict.from_name(
|
|
settings.interrupt_dict_name,
|
|
create_if_missing=True,
|
|
)
|
|
node_output_cache = modal.Dict.from_name(
|
|
settings.node_output_cache_dict_name,
|
|
create_if_missing=True,
|
|
)
|
|
session_bridge_cache = modal.Dict.from_name(
|
|
settings.session_bridge_dict_name,
|
|
create_if_missing=True,
|
|
)
|
|
invocation_records = modal.Dict.from_name(
|
|
settings.invocation_dict_name,
|
|
create_if_missing=True,
|
|
)
|
|
snapshot_profiles = modal.Dict.from_name(
|
|
settings.snapshot_profile_dict_name,
|
|
create_if_missing=True,
|
|
)
|
|
register_cloud_runtime_stores(
|
|
session_bridge_cache=session_bridge_cache,
|
|
invocation_records=invocation_records,
|
|
volume=vol,
|
|
snapshot_profiles=snapshot_profiles,
|
|
node_output_cache=node_output_cache,
|
|
interrupt_flags=interrupt_flags,
|
|
)
|
|
custom_node_packages = _custom_node_runtime_packages(settings.custom_nodes_dir)
|
|
torch_build = _select_remote_torch_build(settings.modal_gpu)
|
|
runtime_identity = build_remote_runtime_identity(
|
|
repo_root=_REPO_ROOT,
|
|
comfyui_root=settings.comfyui_root,
|
|
custom_nodes_dir=settings.custom_nodes_dir,
|
|
settings=settings,
|
|
)
|
|
logger.info(
|
|
"Building Modal runtime fingerprint=%s protocol=%d python=%s.",
|
|
runtime_identity.fingerprint,
|
|
_REMOTE_APP_PROTOCOL_VERSION,
|
|
REMOTE_PYTHON_VERSION,
|
|
)
|
|
logger.info(
|
|
"Selected Modal PyTorch build gpu=%s cuda=%s install_layers=%s.",
|
|
settings.modal_gpu,
|
|
torch_build.cuda_version,
|
|
torch_build.install_layers,
|
|
)
|
|
image = (
|
|
modal.Image.debian_slim(python_version=REMOTE_PYTHON_VERSION)
|
|
.apt_install(*_comfyui_apt_packages())
|
|
.pip_install(*_comfyui_runtime_packages())
|
|
)
|
|
image = _install_custom_node_packages(image, custom_node_packages)
|
|
image = _install_remote_torch_build(image, torch_build)
|
|
image = _install_remote_accelerator_packages(image, settings.modal_gpu)
|
|
image = image.env(
|
|
_modal_image_environment(settings, runtime_identity.fingerprint)
|
|
)
|
|
image = image.add_local_dir(
|
|
_REPO_ROOT,
|
|
remote_path="/root/comfyui_modal_sync_repo",
|
|
ignore=_should_ignore_repo_path,
|
|
)
|
|
if settings.comfyui_root is not None and settings.comfyui_root.exists():
|
|
image = image.add_local_dir(
|
|
settings.comfyui_root,
|
|
remote_path=str(_REMOTE_COMFYUI_ROOT),
|
|
ignore=_should_ignore_comfyui_path,
|
|
)
|
|
logger.info(
|
|
"Including local ComfyUI checkout %s in Modal image at %s.",
|
|
settings.comfyui_root,
|
|
_REMOTE_COMFYUI_ROOT,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"No local ComfyUI checkout was discovered; remote Modal execution may fail to import ComfyUI core modules."
|
|
)
|
|
|
|
stager_image = (
|
|
modal.Image.debian_slim(python_version=REMOTE_PYTHON_VERSION)
|
|
.env(
|
|
_model_stager_image_environment(
|
|
settings,
|
|
runtime_identity.fingerprint,
|
|
)
|
|
)
|
|
.pip_install(*_remote_huggingface_packages())
|
|
.run_commands(_remote_huggingface_validation_command())
|
|
.add_local_dir(
|
|
_REPO_ROOT,
|
|
remote_path="/root/comfyui_modal_sync_repo",
|
|
ignore=_should_ignore_repo_path,
|
|
)
|
|
)
|
|
|
|
@app.cls(
|
|
image=stager_image,
|
|
volumes={settings.remote_storage_root: vol},
|
|
secrets=[modal_secret],
|
|
cpu=4.0,
|
|
memory=16384,
|
|
max_containers=1,
|
|
scaledown_window=300,
|
|
timeout=7200,
|
|
)
|
|
@modal.concurrent(max_inputs=1)
|
|
class ModelStager:
|
|
"""Resolve and stage pinned Hugging Face snapshots without consuming GPU time."""
|
|
|
|
def _stage_profiles(
|
|
self,
|
|
model_references: list[str],
|
|
resolved_profiles: Mapping[str, Any] | None = None,
|
|
progress_callback: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Resolve and stage profiles while optionally publishing progress."""
|
|
staged_profiles = resolve_and_stage_model_references(
|
|
model_references,
|
|
settings.remote_storage_root,
|
|
resolved_profiles=resolved_profiles,
|
|
owner_id=f"modal:{os.getpid()}:{time.time_ns()}",
|
|
progress_callback=(
|
|
lambda progress: progress_callback(
|
|
{
|
|
"stage": progress.stage,
|
|
"message": progress.message,
|
|
"value": progress.value,
|
|
"max": progress.maximum,
|
|
"unit": progress.unit,
|
|
"indeterminate": progress.indeterminate,
|
|
"model_reference": progress.model_reference,
|
|
}
|
|
)
|
|
if progress_callback is not None
|
|
else None
|
|
),
|
|
)
|
|
results = [profile.to_dict() for profile in staged_profiles]
|
|
if any(
|
|
result["downloaded"]
|
|
or result["manifest_created"]
|
|
for result in results
|
|
):
|
|
vol.commit()
|
|
else:
|
|
logger.info(
|
|
"Skipping Modal Volume commit because all requested LLM "
|
|
"profiles and weights were already durable."
|
|
)
|
|
logger.info(
|
|
"Modal LLM CPU resolution and staging completed for models=%s.",
|
|
model_references,
|
|
)
|
|
return results
|
|
|
|
@modal.method()
|
|
def stage_profiles(
|
|
self,
|
|
model_references: list[str],
|
|
resolved_profiles: Mapping[str, Any] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Resolve model references, stage snapshots, and return metadata."""
|
|
return self._stage_profiles(model_references, resolved_profiles)
|
|
|
|
@modal.method()
|
|
def stage_profiles_stream(
|
|
self,
|
|
model_references: list[str],
|
|
resolved_profiles: Mapping[str, Any] | None = None,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Stream CPU staging progress and finish with resolved profile data."""
|
|
progress_events: queue.Queue[dict[str, Any]] = queue.Queue()
|
|
results: list[list[dict[str, Any]]] = []
|
|
errors: list[Exception] = []
|
|
|
|
def run_staging() -> None:
|
|
"""Run blocking Hugging Face work while the generator yields events."""
|
|
try:
|
|
results.append(
|
|
self._stage_profiles(
|
|
model_references,
|
|
resolved_profiles,
|
|
progress_events.put,
|
|
)
|
|
)
|
|
except Exception as error:
|
|
errors.append(error)
|
|
|
|
staging_thread = threading.Thread(
|
|
target=run_staging,
|
|
name="modal-llm-model-stager",
|
|
daemon=True,
|
|
)
|
|
staging_thread.start()
|
|
while staging_thread.is_alive() or not progress_events.empty():
|
|
try:
|
|
progress = progress_events.get(timeout=0.25)
|
|
except queue.Empty:
|
|
continue
|
|
yield {"kind": "progress", **progress}
|
|
staging_thread.join()
|
|
if errors:
|
|
raise errors[0]
|
|
yield {"kind": "result", "results": results[0] if results else []}
|
|
|
|
@modal.method()
|
|
def runtime_version(self) -> dict[str, Any]:
|
|
"""Return deployment identity without allocating a GPU container."""
|
|
return {
|
|
"protocol_version": _REMOTE_APP_PROTOCOL_VERSION,
|
|
"app_name": __comfy_modal_app_name__,
|
|
"runtime_fingerprint": os.environ.get(
|
|
"COMFY_MODAL_RUNTIME_FINGERPRINT",
|
|
"",
|
|
),
|
|
"python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
|
|
}
|
|
|
|
@app.cls(
|
|
**_remote_engine_cls_options(
|
|
settings,
|
|
vol,
|
|
image,
|
|
modal_secret,
|
|
llm_compile_cache_vol,
|
|
)
|
|
)
|
|
@modal.concurrent(max_inputs=1)
|
|
class RemoteEngine:
|
|
"""Modal runtime class that executes proxied ComfyUI payloads."""
|
|
|
|
snapshot_profile_key: str = modal.parameter(default="")
|
|
gpu_snapshot_enabled: bool = modal.parameter(default=False)
|
|
worker_affinity_key: str = modal.parameter(default="worker-pool:slot:0")
|
|
|
|
@modal.enter(snap=True)
|
|
def setup_snapshot_state(self) -> None:
|
|
"""Prepare snapshot-friendly runtime state before Modal captures memory."""
|
|
with _timed_phase("remote_engine_setup_snapshot"):
|
|
_prewarm_snapshot_state(
|
|
gpu_snapshot_enabled=bool(self.gpu_snapshot_enabled),
|
|
snapshot_profile_key=self.snapshot_profile_key,
|
|
)
|
|
logger.info(
|
|
"RemoteEngine snapshot setup complete for snapshot_profile_key=%s gpu_snapshot_enabled=%s worker_affinity=%s.",
|
|
self.snapshot_profile_key or None,
|
|
bool(self.gpu_snapshot_enabled),
|
|
self.worker_affinity_key,
|
|
)
|
|
|
|
@modal.enter(snap=False)
|
|
def setup_restored_runtime(self) -> None:
|
|
"""Prepare request-serving runtime state after a fresh boot or snapshot restore."""
|
|
with _timed_phase("remote_engine_setup_restored"):
|
|
_prewarm_restored_runtime(llm_compile_cache_vol)
|
|
logger.info(
|
|
"RemoteEngine restored-runtime setup complete for snapshot_profile_key=%s.",
|
|
self.snapshot_profile_key or None,
|
|
)
|
|
|
|
@modal.method()
|
|
def execute_payload(
|
|
self, payload: dict[str, Any], kwargs_payload: bytes
|
|
) -> bytes:
|
|
"""Execute a proxied node or subgraph inside the Modal container."""
|
|
_observe_remote_workflow_for_llm_mode(payload)
|
|
component_id = payload.get("component_id", "single-node")
|
|
reload_marker = _modal_volume_reload_marker(payload)
|
|
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"),
|
|
):
|
|
_hydrate_missing_payload_volume_paths(vol, payload)
|
|
if _should_reload_modal_volume(payload):
|
|
_reload_modal_volume_for_request(
|
|
vol,
|
|
str(component_id),
|
|
reload_marker=reload_marker,
|
|
payload=payload,
|
|
)
|
|
else:
|
|
_emit_modal_volume_reload_skip(component_id, payload)
|
|
|
|
def execute_once() -> bytes:
|
|
"""Execute the underlying payload once inside this request context."""
|
|
if payload.get("payload_kind") == "canary":
|
|
return _execute_canary_payload(
|
|
payload,
|
|
kwargs_payload,
|
|
cancellation_event=execution_control.cancellation_event,
|
|
interrupt_store=interrupt_flags,
|
|
interrupt_flag_key=execution_control.interrupt_flag_key,
|
|
)
|
|
if payload.get("payload_kind") == "mapped_subgraph":
|
|
custom_nodes_root = _extract_custom_nodes_bundle(
|
|
payload.get("custom_nodes_bundle")
|
|
)
|
|
_ensure_comfy_runtime_initialized(custom_nodes_root)
|
|
hydrated_inputs = deserialize_node_inputs(
|
|
kwargs_payload
|
|
)
|
|
return serialize_node_outputs(
|
|
_execute_mapped_subgraph_payload(
|
|
payload,
|
|
hydrated_inputs,
|
|
custom_nodes_root,
|
|
cancellation_event=execution_control.cancellation_event,
|
|
interrupt_store=interrupt_flags,
|
|
interrupt_flag_key=execution_control.interrupt_flag_key,
|
|
)
|
|
)
|
|
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,
|
|
)
|
|
|
|
compile_miss_checkpoint = _llm_compile_miss_checkpoint(payload)
|
|
result = _execute_with_durable_invocation(
|
|
payload, execute_once
|
|
)
|
|
_commit_actual_llm_compile_cache(
|
|
compile_miss_checkpoint,
|
|
llm_compile_cache_vol,
|
|
)
|
|
return result
|
|
except Exception as exc:
|
|
_maybe_schedule_container_termination_on_error(payload, exc)
|
|
raise
|
|
|
|
@modal.method()
|
|
def warmup_for_request(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Prime the current or a newly started Modal container for one prompt."""
|
|
return _prepare_warm_container_for_request(
|
|
vol,
|
|
payload,
|
|
llm_compile_cache_vol,
|
|
)
|
|
|
|
@modal.method()
|
|
def keepalive_for_local_gap(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Keep this affinity slot active while the workflow executes locally."""
|
|
logger.info(
|
|
"Remote local-gap keepalive prompt=%s component=%s worker_affinity=%s.",
|
|
payload.get("prompt_id"),
|
|
payload.get("component_id"),
|
|
self.worker_affinity_key,
|
|
)
|
|
return {
|
|
"component_id": str(payload.get("component_id") or "modal-keepalive"),
|
|
"task_id": os.getenv("MODAL_TASK_ID"),
|
|
"worker_affinity_key": self.worker_affinity_key,
|
|
}
|
|
|
|
@modal.method()
|
|
def runtime_version(self) -> dict[str, Any]:
|
|
"""Return the deployed runtime identity expected by the local client."""
|
|
return {
|
|
"protocol_version": _REMOTE_APP_PROTOCOL_VERSION,
|
|
"app_name": __comfy_modal_app_name__,
|
|
"runtime_fingerprint": os.environ.get(
|
|
"COMFY_MODAL_RUNTIME_FINGERPRINT",
|
|
"",
|
|
),
|
|
"python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
|
|
"vllm_version": importlib.metadata.version("vllm"),
|
|
}
|
|
|
|
@modal.method()
|
|
def execute_payload_stream(
|
|
self,
|
|
payload: dict[str, Any],
|
|
kwargs_payload: bytes,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Stream progress envelopes and a final serialized result for one payload."""
|
|
_observe_remote_workflow_for_llm_mode(payload)
|
|
component_id = payload.get("component_id", "single-node")
|
|
reload_marker = _modal_volume_reload_marker(payload)
|
|
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"),
|
|
):
|
|
_hydrate_missing_payload_volume_paths(vol, payload)
|
|
if _should_reload_modal_volume(payload):
|
|
_reload_modal_volume_for_request(
|
|
vol,
|
|
str(component_id),
|
|
reload_marker=reload_marker,
|
|
payload=payload,
|
|
)
|
|
else:
|
|
_emit_modal_volume_reload_skip(component_id, payload)
|
|
compile_miss_checkpoint = _llm_compile_miss_checkpoint(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,
|
|
)
|
|
_commit_actual_llm_compile_cache(
|
|
compile_miss_checkpoint,
|
|
llm_compile_cache_vol,
|
|
)
|
|
except Exception as exc:
|
|
_maybe_schedule_container_termination_on_error(payload, exc)
|
|
raise
|
|
|
|
else:
|
|
app = None
|
|
RemoteEngine = None
|