"""Cloud snapshot, loader, LLM, and warm-container preparation.""" from __future__ import annotations from concurrent.futures import ThreadPoolExecutor from contextlib import AbstractContextManager import copy from dataclasses import dataclass import hashlib import importlib import json import logging import os from pathlib import Path import threading import time from typing import Any, Callable, Mapping try: from .cloud_comfy_bootstrap import ( _ensure_comfy_runtime_initialized, _ensure_comfyui_support_packages, _extract_custom_nodes_bundle, _load_execution_module, _register_custom_nodes_root, ) from .cloud_image_env import _REMOTE_LLM_COMPILE_CACHE_ROOT from .cloud_prompt_execution import _execute_subgraph_prompt, _is_link from .cloud_session_bridge import _load_loader_snapshot_profile from .cloud_volume_reload import ( _emit_modal_volume_reload_skip, _hydrate_missing_payload_volume_paths, _is_modal_volume_open_files_error, _modal_volume_reload_marker, _reload_modal_volume_for_request, _should_reload_modal_volume, ) from .settings import get_settings except ImportError: # pragma: no cover - flat Modal-container import. from cloud_comfy_bootstrap import ( _ensure_comfy_runtime_initialized, _ensure_comfyui_support_packages, _extract_custom_nodes_bundle, _load_execution_module, _register_custom_nodes_root, ) from cloud_image_env import _REMOTE_LLM_COMPILE_CACHE_ROOT from cloud_prompt_execution import _execute_subgraph_prompt, _is_link from cloud_session_bridge import _load_loader_snapshot_profile from cloud_volume_reload import ( _emit_modal_volume_reload_skip, _hydrate_missing_payload_volume_paths, _is_modal_volume_open_files_error, _modal_volume_reload_marker, _reload_modal_volume_for_request, _should_reload_modal_volume, ) from settings import get_settings logger = logging.getLogger(__name__) _LOADER_PREWARM_PLAN_KEYS_LOCK = threading.Lock() _LOADER_PREWARM_PLAN_KEYS: set[str] = set() _LLM_PREWARM_PLAN_KEYS_LOCK = threading.Lock() _LLM_PREWARM_PLAN_KEYS: set[str] = set() @dataclass(frozen=True) class CloudPrewarmHooks: """Cloud logging callbacks supplied by the stable entrypoint.""" emit_cloud_info: Callable[..., None] timed_phase: Callable[..., AbstractContextManager[None]] _PREWARM_HOOKS: CloudPrewarmHooks | None = None def configure_cloud_prewarm_hooks(hooks: CloudPrewarmHooks) -> None: """Install prewarm callbacks without importing upward.""" global _PREWARM_HOOKS _PREWARM_HOOKS = hooks def _prewarm_hooks() -> CloudPrewarmHooks: """Return configured callbacks or fail on invalid import order.""" if _PREWARM_HOOKS is None: raise RuntimeError("Cloud prewarm hooks have not been configured.") return _PREWARM_HOOKS def _emit_cloud_info(message: str, *args: Any) -> None: """Delegate timestamped cloud logging to the stable entrypoint.""" _prewarm_hooks().emit_cloud_info(message, *args) def _timed_phase(phase: str, **fields: Any) -> AbstractContextManager[None]: """Delegate phase timing to the stable entrypoint.""" return _prewarm_hooks().timed_phase(phase, **fields) def _prewarm_snapshot_state( *, gpu_snapshot_enabled: bool, snapshot_profile_key: str = "", ) -> None: """Run snapshot-safe initialization before Modal captures a memory snapshot.""" with _timed_phase( "prewarm_snapshot_state", gpu_snapshot=gpu_snapshot_enabled, snapshot_profile=snapshot_profile_key or None, ): _ensure_comfyui_support_packages() normalized_snapshot_profile_key = snapshot_profile_key.strip() if gpu_snapshot_enabled and normalized_snapshot_profile_key: _ensure_comfy_runtime_initialized(None) _load_execution_module() loader_prewarm_plans = _load_loader_snapshot_profile( normalized_snapshot_profile_key ) if loader_prewarm_plans: _execute_loader_prewarm_plans( component_id=f"snapshot-profile:{normalized_snapshot_profile_key}", loader_prewarm_plans=loader_prewarm_plans, custom_nodes_root=None, ) _emit_cloud_info( "Completed GPU-snapshot ComfyUI prewarm before snapshot capture." ) return if gpu_snapshot_enabled: _emit_cloud_info( "Skipping GPU-snapshot ComfyUI prewarm before snapshot capture because no snapshot profile was provided." ) else: _emit_cloud_info( "Skipping full ComfyUI runtime prewarm during CPU-only snapshot to avoid accidental CUDA initialization." ) def _reload_compile_cache_volume(volume: Any | None) -> bool: """Refresh persistent compiler artifacts before a runtime opens its caches.""" if volume is None: return False reload_method = getattr(volume, "reload", None) if not callable(reload_method): logger.warning("Modal compile-cache Volume does not expose reload().") return False with _timed_phase("llm_compile_cache_reload"): try: reload_method() except RuntimeError as exc: if _is_modal_volume_open_files_error(exc): _log_compile_cache_memory_maps() raise return True def _mapped_process_files_under( volume_root: Path, *, proc_root: Path = Path("/proc"), ) -> tuple[tuple[int, str], ...]: """Return process ids and files memory-mapped beneath one filesystem root.""" try: resolved_root = volume_root.resolve(strict=True) process_directories = tuple(proc_root.iterdir()) except OSError: return () mapped_files: set[tuple[int, str]] = set() for process_directory in process_directories: if not process_directory.name.isdecimal(): continue try: maps_text = (process_directory / "maps").read_text( encoding="utf-8", errors="replace", ) except OSError: continue for line in maps_text.splitlines(): fields = line.split(maxsplit=5) if len(fields) != 6 or not fields[5].startswith("/"): continue mapped_path = fields[5].removesuffix(" (deleted)") try: Path(mapped_path).relative_to(resolved_root) except ValueError: continue mapped_files.add((int(process_directory.name), mapped_path)) return tuple(sorted(mapped_files)) def _log_compile_cache_memory_maps() -> None: """Log native mappings that explain a busy compile-cache Volume reload.""" mapped_files = _mapped_process_files_under(_REMOTE_LLM_COMPILE_CACHE_ROOT) if not mapped_files: logger.warning( "Modal compile-cache Volume reload reported open files, but no " "memory-mapped cache files were visible in /proc." ) return logger.error( "Modal compile-cache Volume reload is blocked by %d memory-mapped " "native cache file(s): %s", len(mapped_files), [ {"pid": process_id, "path": mapped_path} for process_id, mapped_path in mapped_files[:8] ], ) def _prewarm_restored_runtime(compile_cache_volume: Any | None = None) -> None: """Run post-restore initialization that should be ready before serving requests.""" with _timed_phase("prewarm_restored_runtime"): _reload_compile_cache_volume(compile_cache_volume) _ensure_comfy_runtime_initialized(None) _load_execution_module() def _prepare_warm_container_for_request( volume: Any, payload: dict[str, Any], compile_cache_volume: Any | None = None, ) -> dict[str, Any]: """Prime one RemoteEngine container for a request before the first real execution payload arrives.""" component_id = str(payload.get("component_id") or "modal-warmup") reload_marker = _modal_volume_reload_marker(payload) _hydrate_missing_payload_volume_paths(volume, payload) needs_volume_reload = _should_reload_modal_volume(payload) with _timed_phase("remote_engine_warmup", component=component_id): if needs_volume_reload: _reload_modal_volume_for_request( volume, component_id, reload_marker=reload_marker, payload=payload, ) else: _emit_modal_volume_reload_skip(component_id, payload) custom_nodes_bundle = payload.get("custom_nodes_bundle") custom_nodes_root: Path | None = None if isinstance(custom_nodes_bundle, str) and custom_nodes_bundle.strip(): custom_nodes_root = _extract_custom_nodes_bundle(custom_nodes_bundle) if custom_nodes_root is not None: _register_custom_nodes_root(custom_nodes_root) loader_prewarm_plans = payload.get("loader_prewarm_plans") if isinstance(loader_prewarm_plans, list) and loader_prewarm_plans: _execute_loader_prewarm_plans( component_id=component_id, loader_prewarm_plans=loader_prewarm_plans, custom_nodes_root=custom_nodes_root, ) llm_prewarm_plans = payload.get("llm_prewarm_plans") llm_prewarm_results: list[dict[str, Any]] = [] if isinstance(llm_prewarm_plans, list) and llm_prewarm_plans: llm_prewarm_results = _execute_llm_prewarm_plans( component_id=component_id, prompt_id=( str(payload["prompt_id"]) if payload.get("prompt_id") is not None else None ), llm_prewarm_plans=llm_prewarm_plans, compile_cache_volume=compile_cache_volume, ) return { "component_id": component_id, "task_id": os.getenv("MODAL_TASK_ID"), "warmup_slot_index": ( int(payload["warmup_slot_index"]) if payload.get("warmup_slot_index") is not None else None ), "reloaded_volume": needs_volume_reload, "llm_prewarm_results": llm_prewarm_results, } def _loader_prewarm_plan_key(plan: Mapping[str, Any]) -> str | None: """Return the stable worker-local dedupe key for one loader prewarm plan.""" signature = plan.get("signature") if signature is None: return None normalized_signature = str(signature).strip() return normalized_signature or None def _build_loader_prewarm_payload( *, component_id: str, plan_index: int, plan: Mapping[str, Any], ) -> dict[str, Any]: """Build one synthetic single-node subgraph payload for loader warmup.""" plan_node_id = str(plan.get("node_id") or f"loader-{plan_index}") prompt_id = plan.get("prompt_id") return { "payload_kind": "subgraph", "component_id": f"{component_id}::loader-prewarm:{plan_node_id}", "prompt_id": (str(prompt_id) if prompt_id is not None else None), "component_node_ids": [plan_node_id], "subgraph_prompt": copy.deepcopy(dict(plan["subgraph_prompt"])), "boundary_inputs": [], "boundary_outputs": [], "execute_node_ids": list(plan.get("execute_node_ids") or [plan_node_id]), "extra_data": {}, } def _execute_loader_prewarm_plans( *, component_id: str, loader_prewarm_plans: list[dict[str, Any]], custom_nodes_root: Path | None, ) -> None: """Execute synthetic one-node loader workflows so fresh workers preload heavyweight models.""" if not get_settings().enable_loader_prewarm: return _ensure_comfy_runtime_initialized(custom_nodes_root) executable_plans: list[tuple[int, Mapping[str, Any], str | None]] = [] skipped_plan_count = 0 for plan_index, plan in enumerate(loader_prewarm_plans): if not isinstance(plan, Mapping): continue plan_key = _loader_prewarm_plan_key(plan) if plan_key is not None: with _LOADER_PREWARM_PLAN_KEYS_LOCK: if plan_key in _LOADER_PREWARM_PLAN_KEYS: skipped_plan_count += 1 continue _LOADER_PREWARM_PLAN_KEYS.add(plan_key) executable_plans.append((plan_index, plan, plan_key)) def execute_plan( plan_entry: tuple[int, Mapping[str, Any], str | None] ) -> None: """Execute one reserved loader plan and make failures retryable.""" plan_index, plan, plan_key = plan_entry started_at = time.perf_counter() try: _execute_subgraph_prompt( _build_loader_prewarm_payload( component_id=component_id, plan_index=plan_index, plan=plan, ), hydrated_inputs={}, custom_nodes_root=custom_nodes_root, ) except Exception: if plan_key is not None: with _LOADER_PREWARM_PLAN_KEYS_LOCK: _LOADER_PREWARM_PLAN_KEYS.discard(plan_key) raise logger.info( "Completed loader prewarm component=%s class_type=%s plan_index=%d elapsed_seconds=%.3f.", component_id, plan.get("class_type"), plan_index, time.perf_counter() - started_at, ) worker_count = min( len(executable_plans), max(1, int(get_settings().loader_prewarm_workers)), ) if worker_count > 1: logger.info( "Running %d independent loader prewarms with bounded concurrency=%d component=%s.", len(executable_plans), worker_count, component_id, ) with ThreadPoolExecutor( max_workers=worker_count, thread_name_prefix="modal-loader-prewarm", ) as executor: futures = [ executor.submit(execute_plan, plan_entry) for plan_entry in executable_plans ] for future in futures: future.result() else: for plan_entry in executable_plans: execute_plan(plan_entry) executed_plan_count = len(executable_plans) if executed_plan_count or skipped_plan_count: logger.info( "Warm container loader prewarm finished for component=%s executed=%d skipped=%d.", component_id, executed_plan_count, skipped_plan_count, ) def _llm_prewarm_model_profile(plan: Mapping[str, Any]) -> str: """Return the staged model profile from one rewritten LLM warmup plan.""" prompt_node = plan.get("prompt_node") if isinstance(prompt_node, Mapping): inputs = prompt_node.get("inputs") if isinstance(inputs, Mapping): model_profile = inputs.get("model_profile") if isinstance(model_profile, str) and model_profile.strip(): return model_profile.strip() model_profile = plan.get("model_profile") if not isinstance(model_profile, str) or not model_profile.strip(): raise ValueError("LLM prewarm plan requires a fixed model_profile.") return model_profile.strip() def _llm_compile_manifest_path(signature: str) -> Path: """Return the content-addressed completion marker for one JIT warmup plan.""" cache_root = Path( os.getenv("TRITON_CACHE_DIR", str(_REMOTE_LLM_COMPILE_CACHE_ROOT)) ).parent return cache_root / "manifests" / f"{signature}.json" def _write_llm_compile_manifest( *, signature: str, model_profile: str, result: Mapping[str, Any], ) -> Path: """Atomically publish successful representative-warmup metadata.""" manifest_path = _llm_compile_manifest_path(signature) manifest_path.parent.mkdir(parents=True, exist_ok=True) temporary_path = manifest_path.with_suffix(f".{os.getpid()}.tmp") temporary_path.write_text( json.dumps( { "signature": signature, "model_profile": model_profile, "runtime_fingerprint": os.getenv( "COMFY_MODAL_RUNTIME_FINGERPRINT", "" ), "completed_at": time.time(), "result": dict(result), }, sort_keys=True, ), encoding="utf-8", ) temporary_path.replace(manifest_path) return manifest_path def _execute_llm_prewarm_plans( *, component_id: str, prompt_id: str | None, llm_prewarm_plans: list[dict[str, Any]], compile_cache_volume: Any | None, ) -> list[dict[str, Any]]: """Load resident LLMs, exercise representative shapes, and commit JIT caches.""" from modal_llm_runtime import prewarm_modal_llm_profile results: list[dict[str, Any]] = [] for plan in llm_prewarm_plans: if not isinstance(plan, Mapping): continue plan_signature = str(plan.get("signature") or "").strip() if not plan_signature: raise ValueError("LLM prewarm plan requires a stable signature.") model_profile = _llm_prewarm_model_profile(plan) signature = hashlib.sha256( json.dumps( { "model_profile": model_profile, "plan_signature": plan_signature, "runtime_fingerprint": os.getenv( "COMFY_MODAL_RUNTIME_FINGERPRINT", "" ), }, sort_keys=True, ).encode("utf-8") ).hexdigest() manifest_path = _llm_compile_manifest_path(signature) representative_request_count = ( 1 if manifest_path.exists() else max(1, int(plan.get("representative_request_count") or 3)) ) with _LLM_PREWARM_PLAN_KEYS_LOCK: already_resident = signature in _LLM_PREWARM_PLAN_KEYS _LLM_PREWARM_PLAN_KEYS.add(signature) if already_resident: logger.info( "Skipping duplicate resident LLM prewarm profile=%s component=%s.", model_profile, component_id, ) continue try: compile_checkpoint = _LLMCompileMissCheckpoint( profiles=(model_profile,), signal_size=_triton_compile_miss_signal_size(), listener_engine_pids=_triton_compile_listener_engine_pids(), ) with _timed_phase( "llm_representative_prewarm", component=component_id, profile=model_profile, requests=representative_request_count, ): result = prewarm_modal_llm_profile( model_profile=model_profile, representative_request_count=representative_request_count, workflow_execution_id=prompt_id, ) manifest_path = _write_llm_compile_manifest( signature=signature, model_profile=model_profile, result=result, ) compile_cache_committed = _commit_actual_llm_compile_cache( compile_checkpoint, compile_cache_volume, ) results.append( { **result, "manifest_path": str(manifest_path), "manifest_cache_hit": representative_request_count == 1, "compile_cache_committed": compile_cache_committed, } ) except Exception: with _LLM_PREWARM_PLAN_KEYS_LOCK: _LLM_PREWARM_PLAN_KEYS.discard(signature) raise return results def _llm_profiles_in_payload(payload: Mapping[str, Any]) -> tuple[str, ...]: """Collect ModalLLM profiles in the executable subgraph dependency closure.""" if payload.get("payload_kind") not in {"subgraph", "mapped_subgraph"}: return () subgraph_prompt = payload.get("subgraph_prompt") execute_node_ids = payload.get("execute_node_ids") if not isinstance(subgraph_prompt, Mapping) or not isinstance( execute_node_ids, (list, tuple) ): return () prompt = {str(node_id): node for node_id, node in subgraph_prompt.items()} profiles: set[str] = set() visited: set[str] = set() pending = [str(node_id) for node_id in execute_node_ids] while pending: node_id = pending.pop() if node_id in visited: continue visited.add(node_id) prompt_node = prompt.get(node_id) if not isinstance(prompt_node, Mapping): continue inputs = prompt_node.get("inputs") if prompt_node.get("class_type") == "ModalLLM" and isinstance( inputs, Mapping ): profile = inputs.get("model_profile") if isinstance(profile, str) and profile.strip(): profiles.add(profile.strip()) if not isinstance(inputs, Mapping): continue for input_value in inputs.values(): if _is_link(input_value): pending.append(str(input_value[0])) return tuple(sorted(profiles)) @dataclass(frozen=True) class _LLMCompileMissCheckpoint: """Capture the genuine Triton miss signal before one LLM subgraph executes.""" profiles: tuple[str, ...] signal_size: int listener_engine_pids: tuple[int, ...] def _triton_compile_miss_signal_size() -> int: """Read the EngineCore compile-miss signal shared through container storage.""" runtime_module = importlib.import_module("modal_llm_runtime") signal_reader = getattr(runtime_module, "triton_compile_miss_signal_size", None) if not callable(signal_reader): raise RuntimeError( "Modal LLM runtime does not expose triton_compile_miss_signal_size()." ) return int(signal_reader()) def _triton_compile_listener_engine_pids() -> tuple[int, ...]: """Return live EngineCore processes with cache-aware Triton telemetry.""" runtime_module = importlib.import_module("modal_llm_runtime") listener_reader = getattr( runtime_module, "triton_compile_listener_engine_pids", None, ) if not callable(listener_reader): raise RuntimeError( "Modal LLM runtime does not expose " "triton_compile_listener_engine_pids()." ) return tuple(int(pid) for pid in listener_reader()) def _llm_compile_miss_checkpoint( payload: Mapping[str, Any], ) -> _LLMCompileMissCheckpoint | None: """Capture the current miss signal for an executable ModalLLM subgraph.""" profiles = _llm_profiles_in_payload(payload) if not profiles: return None return _LLMCompileMissCheckpoint( profiles=profiles, signal_size=_triton_compile_miss_signal_size(), listener_engine_pids=_triton_compile_listener_engine_pids(), ) def _commit_actual_llm_compile_cache( checkpoint: _LLMCompileMissCheckpoint | None, compile_cache_volume: Any | None, ) -> bool: """Commit the compile-cache Volume after a genuine Triton disk-cache miss.""" if checkpoint is None or compile_cache_volume is None: return False listener_engine_pids = _triton_compile_listener_engine_pids() if not listener_engine_pids: logger.warning( "Skipping LLM compile-cache commit because no live vLLM EngineCore " "reported the cache-aware Triton listener profiles=%s " "listener_pids_before=%s.", checkpoint.profiles, checkpoint.listener_engine_pids, ) return False signal_size = _triton_compile_miss_signal_size() if signal_size < checkpoint.signal_size: raise RuntimeError( "Triton compile-miss signal shrank during LLM execution: " f"before={checkpoint.signal_size} after={signal_size}." ) if signal_size == checkpoint.signal_size: logger.info( "Skipping LLM compile-cache commit because every Triton lookup hit " "the persistent cache profiles=%s signal_size=%d listener_pids=%s.", checkpoint.profiles, signal_size, listener_engine_pids, ) return False commit_method = getattr(compile_cache_volume, "commit", None) if not callable(commit_method): raise RuntimeError("Modal compile-cache Volume does not expose commit().") with _timed_phase( "llm_actual_compile_cache_commit", profiles=checkpoint.profiles, miss_signal_bytes=signal_size - checkpoint.signal_size, ): commit_method() return True