"""Modal warmup, snapshot-profile, and local-gap keepalive orchestration.""" from __future__ import annotations import asyncio from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError as FutureTimeoutError from contextlib import nullcontext import copy from dataclasses import dataclass, field import hashlib import json import logging import queue import statistics import threading import time from typing import Any, Callable, Iterator, Mapping from ..settings import get_settings, modal_deployment_app_name from .local_execution import _is_link, _resolve_required_subgraph_nodes from .modal_deployment import ( ModalRemoteInvocationError, _auto_deploy_modal_app, _call_modal_method, _component_pool_slot_index, _is_missing_modal_deployment_error, _load_modal_cloud_module, _lookup_deployed_remote_engine, _lookup_protocol_current_remote_engine, _modal_cloud_settings_override, _modal_lookup_error_types, _remote_worker_affinity_key, _settings_for_payload, ) logger = logging.getLogger(__name__) try: import modal # type: ignore except ModuleNotFoundError: # pragma: no cover - local fallback environments. modal = None _PROMPT_WARMUP_STATES_LOCK = threading.Lock() _PROMPT_WARMUP_STATES: dict[str, "_PromptWarmupState"] = {} _PROMPT_WARMUP_STATE_ORDER: queue.SimpleQueue[str] | None = None _PROMPT_WARMUP_STATE_CACHE_LIMIT = 256 _LOCAL_GAP_KEEPALIVES_LOCK = threading.Lock() _LOCAL_GAP_KEEPALIVES: dict[tuple[str, int], "_LocalGapKeepaliveState"] = {} _SNAPSHOT_PROFILE_RECORDS_LOCK = threading.Lock() _SNAPSHOT_PROFILE_RECORDS: dict[str, dict[str, Any]] = {} _ROOT_LOADER_PREWARM_CLASS_TYPES = frozenset( {"CheckpointLoaderSimple", "UNETLoader", "CLIPLoader", "DualCLIPLoader"} ) _SPECULATIVE_PREWARM_TARGET_KEY = "speculative_remote_prewarm_target" _REMOTE_MODAL_WARMUP_EXECUTOR_LOCK = threading.Lock() _REMOTE_MODAL_WARMUP_EXECUTOR: ThreadPoolExecutor | None = None _REMOTE_MODAL_KEEPALIVE_EXECUTOR_LOCK = threading.Lock() _REMOTE_MODAL_KEEPALIVE_EXECUTOR: ThreadPoolExecutor | None = None @dataclass(frozen=True) class ModalWarmupHooks: """Callbacks supplied by the host orchestrator for adjacent concerns.""" ensure_llm_profiles_staged: Callable[[dict[str, Any], str], None] mapped_execution_parallelism: Callable[[int], int] _WARMUP_HOOKS: ModalWarmupHooks | None = None def configure_modal_warmup_hooks(hooks: ModalWarmupHooks) -> None: """Install host callbacks without importing upward into the orchestrator.""" global _WARMUP_HOOKS _WARMUP_HOOKS = hooks def _warmup_hooks() -> ModalWarmupHooks: """Return configured host callbacks or fail with a clear import-order error.""" if _WARMUP_HOOKS is None: raise RuntimeError("Modal warmup hooks have not been configured.") return _WARMUP_HOOKS def _ensure_llm_profiles_staged( payload: dict[str, Any], deployment_app_name: str ) -> None: """Delegate LLM staging through the injected host callback.""" _warmup_hooks().ensure_llm_profiles_staged(payload, deployment_app_name) def _remote_modal_call_worker_count() -> int: """Return the number of local threads reserved for blocking Modal calls.""" return max(1, int(get_settings().max_inflight_calls)) def _remote_modal_warmup_executor() -> ThreadPoolExecutor: """Return the lazily constructed executor for Modal warmup calls.""" global _REMOTE_MODAL_WARMUP_EXECUTOR with _REMOTE_MODAL_WARMUP_EXECUTOR_LOCK: if _REMOTE_MODAL_WARMUP_EXECUTOR is None: _REMOTE_MODAL_WARMUP_EXECUTOR = ThreadPoolExecutor( max_workers=_remote_modal_call_worker_count() ) return _REMOTE_MODAL_WARMUP_EXECUTOR def _remote_modal_keepalive_executor() -> ThreadPoolExecutor: """Return the lazily constructed executor for Modal keepalive calls.""" global _REMOTE_MODAL_KEEPALIVE_EXECUTOR with _REMOTE_MODAL_KEEPALIVE_EXECUTOR_LOCK: if _REMOTE_MODAL_KEEPALIVE_EXECUTOR is None: _REMOTE_MODAL_KEEPALIVE_EXECUTOR = ThreadPoolExecutor( max_workers=_remote_modal_call_worker_count() ) return _REMOTE_MODAL_KEEPALIVE_EXECUTOR @dataclass class _PromptWarmupState: """Track proactive warmup state for one local prompt.""" scheduled_slots: set[int] = field(default_factory=set) exact_component_parallelism: dict[str, int] = field(default_factory=dict) slot_futures: dict[int, Future[Any]] = field(default_factory=dict) scheduled_speculative_affinities: set[str] = field(default_factory=set) speculative_affinity_futures: dict[str, Future[Any]] = field(default_factory=dict) @dataclass class _LocalGapKeepaliveState: """Track one prompt-scoped worker-retention loop during local execution.""" component_id: str stop_event: threading.Event future: Future[Any] def _mapped_execution_parallelism(item_count: int) -> int: """Delegate mapped worker-width calculation through the injected callback.""" return _warmup_hooks().mapped_execution_parallelism(item_count) def _warmup_prompt_id(warmup_request: dict[str, Any]) -> str | None: """Return the prompt id that scopes proactive warmup state.""" prompt_id = warmup_request.get("prompt_id") if prompt_id is None: return None normalized_prompt_id = str(prompt_id).strip() return normalized_prompt_id or None def _ensure_prompt_warmup_state(prompt_id: str) -> _PromptWarmupState: """Return the cached warmup state bucket for one prompt.""" global _PROMPT_WARMUP_STATE_ORDER state = _PROMPT_WARMUP_STATES.get(prompt_id) if state is not None: return state state = _PromptWarmupState() _PROMPT_WARMUP_STATES[prompt_id] = state if _PROMPT_WARMUP_STATE_ORDER is None: _PROMPT_WARMUP_STATE_ORDER = queue.SimpleQueue() _PROMPT_WARMUP_STATE_ORDER.put(prompt_id) while len(_PROMPT_WARMUP_STATES) > _PROMPT_WARMUP_STATE_CACHE_LIMIT: expired_prompt_id = _PROMPT_WARMUP_STATE_ORDER.get() _PROMPT_WARMUP_STATES.pop(expired_prompt_id, None) return state def _clamp_prompt_warmup_target(warmup_target: int) -> int: """Clamp one proactive warmup target to the configured Modal container cap.""" normalized_target = max(0, int(warmup_target)) max_containers = get_settings().max_containers if max_containers is not None: return min(normalized_target, max_containers) return normalized_target def _iter_loader_prewarm_prompt_payloads( payload: dict[str, Any] ) -> Iterator[dict[str, Any]]: """Yield subgraph-like payload fragments that may contain root loader nodes.""" split_proxy_payloads = payload.get("split_proxy_payloads") if isinstance(split_proxy_payloads, dict): for phase_payload in split_proxy_payloads.values(): if isinstance(phase_payload, dict): yield phase_payload return if isinstance(split_proxy_payloads, list): for phase_payload in split_proxy_payloads: if isinstance(phase_payload, dict): yield phase_payload return if isinstance(payload.get("subgraph_prompt"), dict): yield payload def _is_root_literal_loader_node(prompt_node: Mapping[str, Any]) -> bool: """Return whether one prompt node is a supported loader with only literal inputs.""" class_type = str(prompt_node.get("class_type") or "") if class_type not in _ROOT_LOADER_PREWARM_CLASS_TYPES: return False inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): return False return not any(_is_link(input_value) for input_value in inputs.values()) def _loader_prewarm_plan_signature(class_type: str, inputs: Mapping[str, Any]) -> str: """Return a stable signature for one synthetic loader prewarm plan.""" return json.dumps( { "class_type": class_type, "inputs": copy.deepcopy(dict(inputs)), }, sort_keys=True, default=str, ) def _loader_snapshot_profile_key(loader_prewarm_plans: list[dict[str, Any]]) -> str: """Return the stable snapshot-profile key for one set of loader prewarm plans.""" if not loader_prewarm_plans: return "" profile_payload = { "plan_signatures": sorted( str(plan.get("signature") or "") for plan in loader_prewarm_plans if str(plan.get("signature") or "") ) } if not profile_payload["plan_signatures"]: return "" profile_digest = hashlib.sha256( json.dumps(profile_payload, sort_keys=True).encode("utf-8") ).hexdigest() return f"loader-profile:{profile_digest}" def _normalize_loader_snapshot_profile_record( loader_prewarm_plans: list[dict[str, Any]], ) -> dict[str, Any] | None: """Return the normalized snapshot-profile record for one loader-plan set.""" snapshot_profile_key = _loader_snapshot_profile_key(loader_prewarm_plans) if not snapshot_profile_key: return None return { "snapshot_profile_key": snapshot_profile_key, "loader_prewarm_plans": copy.deepcopy(loader_prewarm_plans), "snapshot_policy": { "selected_variant": None, "measurements": {"snapshot": [], "direct": []}, }, } def _store_loader_snapshot_profile(loader_prewarm_plans: list[dict[str, Any]]) -> str: """Persist one loader snapshot profile so `snap=True` can retrieve it later.""" settings = get_settings() if not settings.enable_gpu_memory_snapshot or modal is None: return "" normalized_record = _normalize_loader_snapshot_profile_record(loader_prewarm_plans) if normalized_record is None: return "" snapshot_profile_key = str(normalized_record["snapshot_profile_key"]) with _SNAPSHOT_PROFILE_RECORDS_LOCK: cached_record = _SNAPSHOT_PROFILE_RECORDS.get(snapshot_profile_key) if cached_record is not None: return snapshot_profile_key snapshot_profiles = modal.Dict.from_name( settings.snapshot_profile_dict_name, create_if_missing=True, ) stored_record = snapshot_profiles.get(snapshot_profile_key) if isinstance(stored_record, dict): merged_record = copy.deepcopy(stored_record) merged_record["snapshot_profile_key"] = snapshot_profile_key merged_record["loader_prewarm_plans"] = copy.deepcopy( loader_prewarm_plans ) else: merged_record = normalized_record snapshot_profiles[snapshot_profile_key] = merged_record _SNAPSHOT_PROFILE_RECORDS[snapshot_profile_key] = merged_record return snapshot_profile_key def _prepare_snapshot_profile_fields(payload: dict[str, Any]) -> str: """Derive, persist, and attach the loader profile used by one request.""" loader_prewarm_plans = _build_loader_prewarm_plans(payload) stored_profile_key = _store_loader_snapshot_profile(loader_prewarm_plans) existing_profile_key = str(payload.get("snapshot_profile_key") or "").strip() snapshot_profile_key = stored_profile_key or existing_profile_key if snapshot_profile_key: payload["snapshot_profile_key"] = snapshot_profile_key return snapshot_profile_key def _snapshot_profile_store_for_payload(payload: Mapping[str, Any]) -> Any | None: """Return the shared profile store when Modal and snapshot policy are enabled.""" if modal is None: return None settings = _settings_for_payload(payload) return modal.Dict.from_name( settings.snapshot_profile_dict_name, create_if_missing=True, ) def _snapshot_policy_record( payload: Mapping[str, Any], snapshot_profile_key: str ) -> dict[str, Any] | None: """Load one profile record without discarding locally cached measurements.""" with _SNAPSHOT_PROFILE_RECORDS_LOCK: cached = _SNAPSHOT_PROFILE_RECORDS.get(snapshot_profile_key) store = _snapshot_profile_store_for_payload(payload) if store is None: return copy.deepcopy(cached) if cached is not None else None stored = store.get(snapshot_profile_key) if not isinstance(stored, dict): return copy.deepcopy(cached) if cached is not None else None with _SNAPSHOT_PROFILE_RECORDS_LOCK: _SNAPSHOT_PROFILE_RECORDS[snapshot_profile_key] = copy.deepcopy(stored) return copy.deepcopy(stored) def _select_gpu_snapshot_for_profile( payload: dict[str, Any], snapshot_profile_key: str ) -> bool: """Choose snapshot or direct startup from persisted per-profile measurements.""" settings = _settings_for_payload(payload) if not settings.enable_gpu_memory_snapshot or not snapshot_profile_key: payload["gpu_snapshot_variant"] = "direct" return False requested_variant = str(payload.get("gpu_snapshot_variant") or "").strip() if requested_variant in {"snapshot", "direct"}: return requested_variant == "snapshot" record = _snapshot_policy_record(payload, snapshot_profile_key) or {} policy = record.get("snapshot_policy") if not isinstance(policy, Mapping): policy = {} selected_variant = policy.get("selected_variant") if selected_variant not in {"snapshot", "direct"}: measurements = policy.get("measurements") if not isinstance(measurements, Mapping): measurements = {} snapshot_samples = measurements.get("snapshot") direct_samples = measurements.get("direct") snapshot_count = len(snapshot_samples) if isinstance(snapshot_samples, list) else 0 direct_count = len(direct_samples) if isinstance(direct_samples, list) else 0 selected_variant = "snapshot" if snapshot_count <= direct_count else "direct" payload["gpu_snapshot_variant"] = selected_variant logger.info( "Selected Modal startup variant=%s snapshot_profile=%s component=%s.", selected_variant, snapshot_profile_key, payload.get("component_id"), ) return selected_variant == "snapshot" def _record_snapshot_warmup_measurement( payload: Mapping[str, Any], elapsed_seconds: float ) -> None: """Persist one warmup latency and select the faster profile after both variants run.""" snapshot_profile_key = str(payload.get("snapshot_profile_key") or "").strip() variant = str(payload.get("gpu_snapshot_variant") or "").strip() if not snapshot_profile_key or variant not in {"snapshot", "direct"}: return store = _snapshot_profile_store_for_payload(payload) if store is None: return record = _snapshot_policy_record(payload, snapshot_profile_key) or { "snapshot_profile_key": snapshot_profile_key, "loader_prewarm_plans": [], } policy = record.get("snapshot_policy") if not isinstance(policy, dict): policy = {} measurements = policy.get("measurements") if not isinstance(measurements, dict): measurements = {"snapshot": [], "direct": []} for candidate in ("snapshot", "direct"): samples = measurements.get(candidate) if not isinstance(samples, list): measurements[candidate] = [] measurements[variant].append(float(elapsed_seconds)) measurements[variant] = measurements[variant][-5:] if measurements["snapshot"] and measurements["direct"]: policy["selected_variant"] = min( ("snapshot", "direct"), key=lambda candidate: statistics.median(measurements[candidate]), ) else: policy["selected_variant"] = None policy["measurements"] = measurements record["snapshot_policy"] = policy store[snapshot_profile_key] = record with _SNAPSHOT_PROFILE_RECORDS_LOCK: _SNAPSHOT_PROFILE_RECORDS[snapshot_profile_key] = copy.deepcopy(record) logger.info( "Recorded Modal startup measurement snapshot_profile=%s variant=%s elapsed=%.3fs selected=%s.", snapshot_profile_key, variant, elapsed_seconds, policy["selected_variant"], ) def _build_loader_prewarm_plans(payload: dict[str, Any]) -> list[dict[str, Any]]: """Return synthetic one-node subgraph plans for supported root loader nodes.""" if not get_settings().enable_loader_prewarm: return [] prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) seen_signatures: set[str] = set() plans: list[dict[str, Any]] = [] for prompt_payload in _iter_loader_prewarm_prompt_payloads(payload): subgraph_prompt = prompt_payload.get("subgraph_prompt") if not isinstance(subgraph_prompt, Mapping): continue for node_id, prompt_node in subgraph_prompt.items(): if not isinstance(prompt_node, Mapping) or not _is_root_literal_loader_node( prompt_node ): continue class_type = str(prompt_node.get("class_type") or "") inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue signature = _loader_prewarm_plan_signature(class_type, inputs) if signature in seen_signatures: continue seen_signatures.add(signature) plans.append( { "signature": signature, "node_id": str(node_id), "class_type": class_type, "prompt_id": prompt_id, "subgraph_prompt": { str(node_id): copy.deepcopy(dict(prompt_node)), }, "execute_node_ids": [str(node_id)], } ) return plans def _iter_executable_subgraph_nodes( payload: Mapping[str, Any], ) -> Iterator[Mapping[str, Any]]: """Yield only prompt nodes required by this payload's execution targets.""" for prompt_payload in _iter_loader_prewarm_prompt_payloads(dict(payload)): raw_prompt = prompt_payload.get("subgraph_prompt") if not isinstance(raw_prompt, Mapping): continue prompt = { str(node_id): dict(prompt_node) for node_id, prompt_node in raw_prompt.items() if isinstance(prompt_node, Mapping) } execute_node_ids: list[str] = [] for field_name in ( "execute_node_ids", "mapped_execute_node_ids", "static_execute_node_ids", ): field_value = prompt_payload.get(field_name) if isinstance(field_value, (list, tuple)): execute_node_ids.extend(str(node_id) for node_id in field_value) required_node_ids = _resolve_required_subgraph_nodes( prompt, list(dict.fromkeys(execute_node_ids)), ) for node_id in required_node_ids: yield prompt[node_id] def _build_llm_prewarm_plans(payload: Mapping[str, Any]) -> list[dict[str, Any]]: """Return LLM warmups found in executable subgraph dependency closures.""" plans: list[dict[str, Any]] = [] seen_profiles: set[str] = set() for prompt_node in _iter_executable_subgraph_nodes(payload): if prompt_node.get("class_type") != "ModalLLM": continue inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): continue model_profile = inputs.get("model_profile") if not isinstance(model_profile, str) or not model_profile.strip(): continue normalized_profile = model_profile.strip() if normalized_profile in seen_profiles: continue seen_profiles.add(normalized_profile) signature_payload = { "model_profile": normalized_profile, "representative_request_count": 3, } plans.append( { **signature_payload, "signature": hashlib.sha256( json.dumps(signature_payload, sort_keys=True).encode("utf-8") ).hexdigest(), "prompt_node": copy.deepcopy(dict(prompt_node)), } ) return plans def _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any] | None: """Extract the prompt-scoped warmup-relevant fields from one payload.""" settings = _settings_for_payload(payload) loader_prewarm_plans = _build_loader_prewarm_plans(payload) llm_prewarm_plans = _build_llm_prewarm_plans(payload) snapshot_profile_key = _store_loader_snapshot_profile(loader_prewarm_plans) if settings.enable_gpu_memory_snapshot and not snapshot_profile_key and not llm_prewarm_plans: logger.info( "Skipping proactive Modal warmup request for component=%s because GPU snapshots are enabled and no loader snapshot profile was derived.", payload.get("component_id"), ) return None return { "prompt_id": ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ), "component_id": str( payload.get( "mapped_progress_display_node_id", payload.get("component_id", "modal-warmup"), ) ), "modal_gpu": settings.modal_gpu, "remote_worker_affinity_group": payload.get( "remote_worker_affinity_group" ), "remote_local_gap_pool": bool(payload.get("remote_local_gap_pool")), "requires_volume_reload": bool(payload.get("requires_volume_reload", True)), "volume_reload_marker": payload.get("volume_reload_marker"), "uploaded_volume_paths": list(payload.get("uploaded_volume_paths", [])), "custom_nodes_bundle": payload.get("custom_nodes_bundle"), "loader_prewarm_plans": loader_prewarm_plans, "llm_prewarm_plans": llm_prewarm_plans, "snapshot_profile_key": snapshot_profile_key, } def _component_parallelism_metadata(payload: dict[str, Any]) -> dict[str, Any]: """Return the prompt-level parallelism metadata attached by the queue rewrite.""" extra_data = payload.get("extra_data") or {} modal_metadata = extra_data.get("modal") or {} if not isinstance(modal_metadata, dict): return {} return modal_metadata def _prompt_parallelism_target( payload: dict[str, Any], *, exact_component_parallelism: dict[str, int] | None = None, ) -> int: """Return the current best-effort prompt-wide warmup target for one payload.""" metadata = _component_parallelism_metadata(payload) execution_stages = metadata.get("component_execution_stages") if not isinstance(execution_stages, list): return 1 mapped_component_ids = { str(component_id) for component_id in metadata.get("mapped_component_ids", []) if str(component_id) } exact_parallelism = exact_component_parallelism or {} stage_parallelism = 0 for stage in execution_stages: if not isinstance(stage, list): continue current_stage_parallelism = 0 for component_id_value in stage: component_id = str(component_id_value) if component_id in exact_parallelism: current_stage_parallelism += max( 1, int(exact_parallelism[component_id]) ) elif component_id in mapped_component_ids: current_stage_parallelism += 1 else: current_stage_parallelism += 1 stage_parallelism = max(stage_parallelism, current_stage_parallelism) fallback_parallelism = metadata.get("estimated_max_parallel_requests") if stage_parallelism <= 0 and isinstance(fallback_parallelism, int): stage_parallelism = max(stage_parallelism, fallback_parallelism) return _clamp_prompt_warmup_target(max(1, stage_parallelism)) def _register_exact_component_parallelism( payload: dict[str, Any], component_parallelism: int ) -> int: """Record exact mapped-component parallelism and return the refined prompt-wide warmup target.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) component_id = ( str(payload.get("component_id")) if payload.get("component_id") is not None else None ) if not prompt_id or not component_id: return _clamp_prompt_warmup_target(component_parallelism) with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _ensure_prompt_warmup_state(prompt_id) warmup_state.exact_component_parallelism[component_id] = max( 1, int(component_parallelism) ) exact_parallelism = dict(warmup_state.exact_component_parallelism) return _prompt_parallelism_target( payload, exact_component_parallelism=exact_parallelism ) def _warmup_slot_payload( warmup_request: dict[str, Any], slot_index: int ) -> dict[str, Any]: """Return the remote warmup payload for one desired container slot.""" slot_payload = copy.deepcopy(warmup_request) component_id = str( slot_payload.get("component_id") or slot_payload.get("prompt_id") or "modal-warmup" ) slot_payload["component_id"] = f"{component_id}::warmup:{slot_index}" slot_payload["warmup_slot_index"] = int(slot_index) slot_payload["warmup_only"] = True return slot_payload def _prompt_warmup_head_start_seconds() -> float: """Return the bounded delay the local scheduler may wait for exact warmup slots to start.""" return max(0.0, float(get_settings().proactive_warmup_head_start_seconds)) def _invoke_remote_engine_warmup( remote_engine: Any, warmup_request: dict[str, Any] ) -> Any: """Ask one prepared remote engine instance to warm a container for a prompt.""" warmup_method = getattr(remote_engine, "warmup_for_request", None) if warmup_method is None: logger.warning( "Remote warmup method is unavailable for prompt=%s component=%s; skipping proactive warmup.", warmup_request.get("prompt_id"), warmup_request.get("component_id"), ) return None started_at = time.perf_counter() if hasattr(warmup_method, "remote"): result = warmup_method.remote(warmup_request) else: result = warmup_method(warmup_request) _record_snapshot_warmup_measurement( warmup_request, time.perf_counter() - started_at, ) return result def _invoke_remote_engine_warmup_with_recovery( remote_engine: Any, warmup_request: dict[str, Any], ) -> Any: """Retry one warmup call after auto-deploy when a stale deployed handle vanishes.""" lookup_error_types = _modal_lookup_error_types() settings = _settings_for_payload(warmup_request) try: return _invoke_remote_engine_warmup(remote_engine, warmup_request) except lookup_error_types as exc: if not settings.auto_deploy or not _is_missing_modal_deployment_error(exc): raise logger.warning( "Modal warmup invocation failed for component=%s because the deployed app was missing at call time: %s. Recreating the app and retrying.", warmup_request.get("component_id"), exc, ) recovered_remote_engine = _auto_deploy_modal_app(warmup_request, exc) _ensure_llm_profiles_staged( warmup_request, modal_deployment_app_name(settings), ) return _invoke_remote_engine_warmup(recovered_remote_engine, warmup_request) def _invoke_modal_warmup_blocking(warmup_request: dict[str, Any]) -> Any: """Warm one Modal container slot using deployed or ephemeral app state.""" if modal is None: return None lookup_error_types = _modal_lookup_error_types() settings = _settings_for_payload(warmup_request) deployment_app_name = modal_deployment_app_name(settings) if lookup_error_types: try: remote_engine = _lookup_protocol_current_remote_engine(warmup_request) _ensure_llm_profiles_staged(warmup_request, deployment_app_name) return _invoke_remote_engine_warmup_with_recovery( remote_engine, warmup_request ) except lookup_error_types as exc: if settings.auto_deploy: remote_engine = _auto_deploy_modal_app(warmup_request, exc) try: _ensure_llm_profiles_staged(warmup_request, deployment_app_name) return _invoke_remote_engine_warmup_with_recovery( remote_engine, warmup_request, ) except lookup_error_types as retry_exc: exc = retry_exc if not settings.allow_ephemeral_fallback: raise ModalRemoteInvocationError( "Proactive Modal warmup requires a deployed Modal app or a successful first-run auto-deploy. " f"Lookup failed for app={deployment_app_name!r}: {exc}." ) from exc else: remote_engine = _lookup_protocol_current_remote_engine(warmup_request) _ensure_llm_profiles_staged(warmup_request, deployment_app_name) return _invoke_remote_engine_warmup_with_recovery(remote_engine, warmup_request) with _modal_cloud_settings_override(settings): cloud_module = _load_modal_cloud_module() cloud_app = getattr(cloud_module, "app", None) cloud_remote_engine = getattr(cloud_module, "RemoteEngine", None) if cloud_app is None or cloud_remote_engine is None: raise ModalRemoteInvocationError( "Stable Modal cloud entry module did not expose app and RemoteEngine." ) run_context = cloud_app.run() if hasattr(cloud_app, "run") else nullcontext() with run_context: remote_engine = cloud_remote_engine() return _invoke_remote_engine_warmup_with_recovery(remote_engine, warmup_request) def _run_prompt_warmup_slot( prompt_id: str, slot_index: int, warmup_request: dict[str, Any], reason: str, ) -> None: """Execute one proactive warmup slot and release it for retry on failure.""" try: logger.info( "Starting proactive Modal warmup for prompt=%s slot=%d component=%s reason=%s.", prompt_id, slot_index, warmup_request.get("component_id"), reason, ) _invoke_modal_warmup_blocking(_warmup_slot_payload(warmup_request, slot_index)) except Exception: with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is not None: warmup_state.scheduled_slots.discard(slot_index) logger.exception( "Proactive Modal warmup failed for prompt=%s slot=%d component=%s.", prompt_id, slot_index, warmup_request.get("component_id"), ) def _track_prompt_warmup_future( prompt_id: str, slot_index: int, future: Future[Any], ) -> None: """Track one in-flight warmup future so mapped execution can await short exact warmup bursts.""" with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is not None: warmup_state.slot_futures[slot_index] = future def _clear_tracked_future(completed_future: Future[Any]) -> None: """Drop one completed future from the prompt warmup state.""" with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is None: return tracked_future = warmup_state.slot_futures.get(slot_index) if tracked_future is completed_future: warmup_state.slot_futures.pop(slot_index, None) future.add_done_callback(_clear_tracked_future) async def _await_prompt_warmup_slots( prompt_id: str, slot_indices: list[int], timeout_seconds: float, ) -> int: """Wait briefly for any in-flight prompt warmup slots to finish.""" if timeout_seconds <= 0.0 or not slot_indices: return 0 with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is None: return 0 pending_slot_futures = [ warmup_state.slot_futures[slot_index] for slot_index in slot_indices if slot_index in warmup_state.slot_futures and not warmup_state.slot_futures[slot_index].done() ] if not pending_slot_futures: return 0 wrapped_futures = [ asyncio.wrap_future(slot_future) for slot_future in pending_slot_futures ] done_futures, pending_futures = await asyncio.wait( wrapped_futures, timeout=timeout_seconds, return_when=asyncio.ALL_COMPLETED, ) if done_futures: logger.info( "Prompt=%s warmup head-start completed %d/%d slot(s) before mapped dispatch.", prompt_id, len(done_futures), len(wrapped_futures), ) elif pending_futures: logger.info( "Prompt=%s warmup head-start timed out after %.3fs with %d slot(s) still warming.", prompt_id, timeout_seconds, len(pending_futures), ) return len(done_futures) def ensure_remote_warm_capacity( warmup_request: dict[str, Any] | None, *, warmup_target: int, reason: str, ) -> int: """Best-effort background warmup so enough Modal containers are ready for one prompt.""" settings = get_settings() if not settings.enable_proactive_warmup: return 0 if settings.execution_mode == "local" or modal is None: return 0 if warmup_request is None: return 0 prompt_id = _warmup_prompt_id(warmup_request) if prompt_id is None: return 0 clamped_target = _clamp_prompt_warmup_target(warmup_target) if clamped_target <= 0: return 0 with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _ensure_prompt_warmup_state(prompt_id) missing_slots = [ slot_index for slot_index in range(clamped_target) if slot_index not in warmup_state.scheduled_slots ] for slot_index in missing_slots: warmup_state.scheduled_slots.add(slot_index) if not missing_slots: return clamped_target logger.info( "Scheduling proactive Modal warmup for prompt=%s target=%d missing_slots=%s component=%s reason=%s.", prompt_id, clamped_target, missing_slots, warmup_request.get("component_id"), reason, ) for slot_index in missing_slots: future = _remote_modal_warmup_executor().submit( _run_prompt_warmup_slot, prompt_id, slot_index, copy.deepcopy(warmup_request), reason, ) _track_prompt_warmup_future(prompt_id, slot_index, future) return clamped_target def _run_speculative_affinity_prewarm( prompt_id: str, warmup_identity: str, warmup_request: dict[str, Any], reason: str, ) -> None: """Prepare one future affinity worker and make a failed attempt retryable.""" affinity_key = _remote_worker_affinity_key(warmup_request) try: logger.info( "Starting speculative Modal prewarm prompt=%s affinity=%s component=%s reason=%s.", prompt_id, affinity_key, warmup_request.get("component_id"), reason, ) _invoke_modal_warmup_blocking(_warmup_slot_payload(warmup_request, 0)) logger.info( "Completed speculative Modal prewarm prompt=%s affinity=%s component=%s.", prompt_id, affinity_key, warmup_request.get("component_id"), ) except Exception: with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is not None: warmup_state.scheduled_speculative_affinities.discard( warmup_identity ) logger.exception( "Speculative Modal prewarm failed prompt=%s affinity=%s component=%s.", prompt_id, affinity_key, warmup_request.get("component_id"), ) def _schedule_speculative_affinity_prewarm( payload: Mapping[str, Any], *, reason: str, ) -> bool: """Schedule one planner-selected future worker without blocking current execution.""" target_payload = payload.get(_SPECULATIVE_PREWARM_TARGET_KEY) if not isinstance(target_payload, Mapping): return False if ( payload.get("execution_provider") != "modal" or target_payload.get("execution_provider") != "modal" or payload.get("execution_environment_id") != target_payload.get("execution_environment_id") ): logger.debug( "Skipping speculative Modal prewarm across execution environments " "source=%s target=%s.", payload.get("execution_environment_id"), target_payload.get("execution_environment_id"), ) return False normalized_target_payload = copy.deepcopy(dict(target_payload)) settings = _settings_for_payload(normalized_target_payload) if ( not settings.enable_proactive_warmup or settings.execution_mode == "local" or modal is None or (settings.max_containers is not None and settings.max_containers < 2) ): return False warmup_request = _build_prompt_warmup_request(normalized_target_payload) if warmup_request is None: return False prompt_id = _warmup_prompt_id(warmup_request) if prompt_id is None: return False current_affinity_key = _remote_worker_affinity_key(dict(payload)) target_affinity_key = _remote_worker_affinity_key(warmup_request) if target_affinity_key == current_affinity_key: return False snapshot_profile_key = _prepare_snapshot_profile_fields(warmup_request) _select_gpu_snapshot_for_profile(warmup_request, snapshot_profile_key) warmup_identity = _speculative_warmup_identity(warmup_request) with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _ensure_prompt_warmup_state(prompt_id) if warmup_identity in warmup_state.scheduled_speculative_affinities: return False warmup_state.scheduled_speculative_affinities.add(warmup_identity) logger.info( "Scheduling speculative Modal prewarm prompt=%s current_component=%s current_affinity=%s target_component=%s target_affinity=%s reason=%s.", prompt_id, payload.get("component_id"), current_affinity_key, warmup_request.get("component_id"), target_affinity_key, reason, ) future = _remote_modal_warmup_executor().submit( _run_speculative_affinity_prewarm, prompt_id, warmup_identity, copy.deepcopy(warmup_request), reason, ) with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is not None: warmup_state.speculative_affinity_futures[warmup_identity] = future def _clear_speculative_future(completed_future: Future[Any]) -> None: """Drop a completed speculative future while retaining successful dedupe state.""" with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) if warmup_state is None: return tracked_future = warmup_state.speculative_affinity_futures.get( warmup_identity ) if tracked_future is completed_future: warmup_state.speculative_affinity_futures.pop( warmup_identity, None ) future.add_done_callback(_clear_speculative_future) return True def _speculative_warmup_identity(payload: Mapping[str, Any]) -> str: """Return the exact deployed-worker identity shared by warmup and dispatch.""" settings = _settings_for_payload(payload) identity = { "app": modal_deployment_app_name(settings), "gpu": settings.modal_gpu, "affinity": _remote_worker_affinity_key(dict(payload)), "snapshot_profile": str(payload.get("snapshot_profile_key") or ""), "snapshot_variant": str(payload.get("gpu_snapshot_variant") or "direct"), } return json.dumps(identity, sort_keys=True, separators=(",", ":")) def _await_matching_speculative_prewarm( payload: dict[str, Any], cancellation_event: threading.Event | None, ) -> None: """Join an active speculative warmup instead of racing a duplicate container.""" prompt_id = _warmup_prompt_id(payload) if prompt_id is None: return snapshot_profile_key = _prepare_snapshot_profile_fields(payload) _select_gpu_snapshot_for_profile(payload, snapshot_profile_key) warmup_identity = _speculative_warmup_identity(payload) with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _PROMPT_WARMUP_STATES.get(prompt_id) future = ( warmup_state.speculative_affinity_futures.get(warmup_identity) if warmup_state is not None else None ) if future is None or future.done(): return started_at = time.perf_counter() logger.info( "Joining active speculative Modal prewarm before dispatch prompt=%s component=%s identity=%s.", prompt_id, payload.get("component_id"), warmup_identity, ) while not future.done(): if cancellation_event is not None and cancellation_event.is_set(): raise ModalRemoteInvocationError( "Remote dispatch was cancelled while waiting for speculative warmup." ) try: future.result(timeout=0.1) except FutureTimeoutError: continue future.result() logger.info( "Joined speculative Modal prewarm before dispatch prompt=%s component=%s wait_seconds=%.3f.", prompt_id, payload.get("component_id"), time.perf_counter() - started_at, ) def _schedule_post_deploy_runtime_seed(payload: Mapping[str, Any]) -> bool: """Seed the just-deployed profile and expose its future to the first dispatch.""" warmup_request = _build_prompt_warmup_request(copy.deepcopy(dict(payload))) if warmup_request is None: return False prompt_id = _warmup_prompt_id(warmup_request) if prompt_id is None: return False snapshot_profile_key = _prepare_snapshot_profile_fields(warmup_request) _select_gpu_snapshot_for_profile(warmup_request, snapshot_profile_key) warmup_identity = _speculative_warmup_identity(warmup_request) with _PROMPT_WARMUP_STATES_LOCK: warmup_state = _ensure_prompt_warmup_state(prompt_id) if warmup_identity in warmup_state.scheduled_speculative_affinities: return False warmup_state.scheduled_speculative_affinities.add(warmup_identity) future = _remote_modal_warmup_executor().submit( _run_speculative_affinity_prewarm, prompt_id, warmup_identity, copy.deepcopy(warmup_request), "post_deploy_runtime_seed", ) warmup_state.speculative_affinity_futures[warmup_identity] = future def clear_seed_future(completed_future: Future[Any]) -> None: """Forget the completed seed future while retaining successful dedupe state.""" with _PROMPT_WARMUP_STATES_LOCK: current_state = _PROMPT_WARMUP_STATES.get(prompt_id) if current_state is None: return if ( current_state.speculative_affinity_futures.get(warmup_identity) is completed_future ): current_state.speculative_affinity_futures.pop( warmup_identity, None ) future.add_done_callback(clear_seed_future) logger.info( "Scheduled automatic post-deployment Modal runtime seed prompt=%s component=%s identity=%s.", prompt_id, warmup_request.get("component_id"), warmup_identity, ) return True def _local_gap_keepalive_key(payload: Mapping[str, Any]) -> tuple[str, int] | None: """Return the prompt and affinity-slot key for a local-gap payload.""" prompt_id = _warmup_prompt_id(dict(payload)) if prompt_id is None: return None return (prompt_id, _component_pool_slot_index(dict(payload))) def _stop_local_gap_keepalive(payload: Mapping[str, Any], *, reason: str) -> bool: """Stop a previously scheduled affinity-slot keepalive, if one exists.""" keepalive_key = _local_gap_keepalive_key(payload) if keepalive_key is None: return False with _LOCAL_GAP_KEEPALIVES_LOCK: keepalive_state = _LOCAL_GAP_KEEPALIVES.pop(keepalive_key, None) if keepalive_state is None: return False keepalive_state.stop_event.set() logger.info( "Stopping remote local-gap keepalive prompt=%s slot=%d after component=%s reason=%s.", keepalive_key[0], keepalive_key[1], keepalive_state.component_id, reason, ) return True def _invoke_remote_local_gap_keepalive(payload: dict[str, Any]) -> Any: """Send one lightweight activity pulse to the payload's affinity slot.""" remote_engine = _lookup_deployed_remote_engine(payload) keepalive_method = getattr(remote_engine, "keepalive_for_local_gap", None) if keepalive_method is None: raise ModalRemoteInvocationError( "The deployed RemoteEngine does not expose keepalive_for_local_gap." ) return _call_modal_method(keepalive_method, payload) def _run_local_gap_keepalive( payload: dict[str, Any], stop_event: threading.Event, duration_seconds: float, interval_seconds: float, ) -> None: """Pulse one remote affinity slot until local execution finishes or the budget expires.""" deadline = time.monotonic() + duration_seconds keepalive_count = 0 while not stop_event.is_set(): remaining_seconds = deadline - time.monotonic() if remaining_seconds <= 0.0: break if stop_event.wait(min(interval_seconds, remaining_seconds)): break keepalive_count += 1 pulse_started_at = time.perf_counter() _invoke_remote_local_gap_keepalive(payload) logger.info( "Completed remote local-gap keepalive pulse=%d prompt=%s component=%s in %.3fs.", keepalive_count, payload.get("prompt_id"), payload.get("component_id"), time.perf_counter() - pulse_started_at, ) logger.info( "Remote local-gap keepalive loop finished prompt=%s component=%s pulses=%d stopped=%s.", payload.get("prompt_id"), payload.get("component_id"), keepalive_count, stop_event.is_set(), ) def _start_local_gap_keepalive(payload: Mapping[str, Any]) -> bool: """Retain one remote affinity slot while a downstream local gap executes.""" if ( payload.get("execution_provider") != "modal" or not bool(payload.get("keepalive_after_remote_component")) ): return False settings = _settings_for_payload(payload) duration_seconds = max(0.0, float(settings.local_gap_keepalive_seconds)) if duration_seconds <= 0.0 or settings.execution_mode == "local" or modal is None: return False keepalive_key = _local_gap_keepalive_key(payload) if keepalive_key is None: return False _stop_local_gap_keepalive(payload, reason="replacement") stop_event = threading.Event() keepalive_payload = dict(payload) keepalive_payload["warmup_slot_index"] = keepalive_key[1] keepalive_payload["component_id"] = ( f"{payload.get('component_id', 'modal-component')}::local-gap-keepalive" ) interval_seconds = min( duration_seconds, max(1.0, float(settings.local_gap_keepalive_interval_seconds)), ) future = _remote_modal_keepalive_executor().submit( _run_local_gap_keepalive, keepalive_payload, stop_event, duration_seconds, interval_seconds, ) keepalive_state = _LocalGapKeepaliveState( component_id=str(payload.get("component_id") or "modal-component"), stop_event=stop_event, future=future, ) with _LOCAL_GAP_KEEPALIVES_LOCK: _LOCAL_GAP_KEEPALIVES[keepalive_key] = keepalive_state def clear_completed_keepalive(completed_future: Future[Any]) -> None: """Forget a finished keepalive and report any best-effort failure.""" with _LOCAL_GAP_KEEPALIVES_LOCK: current_state = _LOCAL_GAP_KEEPALIVES.get(keepalive_key) if current_state is keepalive_state: _LOCAL_GAP_KEEPALIVES.pop(keepalive_key, None) failure = completed_future.exception() if failure is not None: logger.warning( "Remote local-gap keepalive failed prompt=%s slot=%d component=%s: %s", keepalive_key[0], keepalive_key[1], keepalive_state.component_id, failure, ) future.add_done_callback(clear_completed_keepalive) logger.info( "Started remote local-gap keepalive prompt=%s slot=%d component=%s duration=%.1fs interval=%.1fs.", keepalive_key[0], keepalive_key[1], keepalive_state.component_id, duration_seconds, interval_seconds, ) return True def boost_mapped_component_warmup( payload: dict[str, Any], *, total_items: int, reason: str, ) -> tuple[int, int]: """Record exact mapped fan-out and top up prompt warmup for the resulting lane count.""" parallelism = _mapped_execution_parallelism(total_items) refined_prompt_warmup_target = _register_exact_component_parallelism( payload, parallelism ) ensure_remote_warm_capacity( _build_prompt_warmup_request(payload), warmup_target=refined_prompt_warmup_target, reason=reason, ) logger.info( "Boosted exact mapped warmup for component=%s total_items=%d local_parallelism=%d prompt_warmup_target=%d reason=%s.", payload.get("component_id"), total_items, parallelism, refined_prompt_warmup_target, reason, ) return parallelism, refined_prompt_warmup_target