diff --git a/README.md b/README.md index 9216cd5..fc10723 100644 --- a/README.md +++ b/README.md @@ -211,6 +211,7 @@ Remote runtime behavior: - When a mapped or implicitly batched component finally hydrates its input batch, Modal-Sync upgrades that queue-time estimate with the exact per-item lane count and tops up background warmup again. That keeps early warmup conservative when fan-out size is unknown, but still lets the runtime race ahead of the heavy component once the true item count is available. - When a local `Modal Map Input` feeds a mapped remote component, the prompt rewrite now registers that node as an exact warmup trigger. As soon as the node executes and sees the concrete list or batch size, it starts the exact container top-up before the mapped proxy begins dispatching remote work. - The mapped scheduler now tracks in-flight warmup slots and can give new exact-parallelism warmup requests a short bounded head-start before lane seeding begins. That trims cold-start stalls without blocking indefinitely on proactive warmup. +- Proactive warmup now also scans each remote payload for root heavy loader nodes with only literal inputs, such as `UNETLoader`, `CheckpointLoaderSimple`, `CLIPLoader`, and `DualCLIPLoader`. Fresh workers execute those loaders as synthetic one-node subgraphs during warmup so later real requests can hit the warm loader cache instead of paying the first model load on the critical path. - If `vol.reload()` still reports open files after that bounded retry window, the worker now checks whether the payload's referenced mounted-volume paths are already visible on disk. When they are, Modal-Sync proceeds and records the reload marker instead of failing on a redundant reload. - The worker now logs mounted-volume reload diagnostics when open-file retries start, including the uploaded paths and referenced paths for that payload, so you can see exactly which files were relevant when reload got stuck behind open handles. - Custom-node bundle sync now treats the digest-index record as authoritative for dedupe and the hash-addressed volume path as the payload location. Once a bundle digest is indexed, later runs reuse it without reuploading or probing the volume for marker blobs. @@ -393,6 +394,7 @@ If a remote-marked node depends on a model filename that cannot be resolved to a - `COMFY_MODAL_MAX_CONTAINERS`: Optional upper bound on simultaneously scaled Modal containers. - `COMFY_MODAL_BUFFER_CONTAINERS`: Optional number of spare warm containers Modal should try to keep ready above current load. - `COMFY_MODAL_ENABLE_PROACTIVE_WARMUP`: Start best-effort background warmup RPCs based on queue-time and runtime parallelism estimates. Default: `true`. +- `COMFY_MODAL_ENABLE_LOADER_PREWARM`: During proactive warmup, execute synthetic one-node loader prompts for root literal model-loader nodes so fresh workers preload common models before the first real payload arrives. Default: `true`. - `COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS`: How long mapped execution may wait for newly requested exact warmup slots before it starts seeding lanes anyway. Default: `2.0`. ## Development diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 70b214a..60b24b9 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -78,6 +78,8 @@ _LOADER_CACHE_WRAPPED_CLASSES: set[str] = set() _LOADER_OUTPUT_CACHE: dict[tuple[str, str], tuple[Any, ...]] = {} _LOADER_CACHE_METRICS_LOCK = threading.Lock() _LOADER_CACHE_METRICS: dict[str, int] = {"hit": 0, "miss": 0} +_LOADER_PREWARM_PLAN_KEYS_LOCK = threading.Lock() +_LOADER_PREWARM_PLAN_KEYS: set[str] = set() _NODE_OUTPUT_CACHE_KEY_PREFIX = "NC_" _BOUNDARY_INPUT_SIGNATURES_KEY = "__comfy_modal_boundary_input_signatures__" _NODE_OUTPUT_CACHE_RECORD_VERSION = 1 @@ -1575,6 +1577,26 @@ def _build_clip_loader_cache_key(kwargs: dict[str, Any]) -> str: ) +def _build_dual_clip_loader_cache_key(kwargs: dict[str, Any]) -> str: + """Build a stable cache key for the ComfyUI dual CLIP loader.""" + import folder_paths + + return _serialize_loader_cache_key( + { + "clip_path_1": folder_paths.get_full_path_or_raise( + "text_encoders", + str(kwargs["clip_name1"]), + ), + "clip_path_2": folder_paths.get_full_path_or_raise( + "text_encoders", + str(kwargs["clip_name2"]), + ), + "type": kwargs.get("type"), + "device": kwargs.get("device", "default"), + } + ) + + def _build_vae_loader_cache_key(kwargs: dict[str, Any]) -> str: """Build a stable cache key for the ComfyUI VAE loader.""" return _serialize_loader_cache_key({"vae_name": kwargs.get("vae_name")}) @@ -1644,6 +1666,7 @@ def _install_loader_cache_wrappers() -> None: "CheckpointLoaderSimple": ("load_checkpoint", _build_checkpoint_loader_cache_key), "UNETLoader": ("load_unet", _build_unet_loader_cache_key), "CLIPLoader": ("load_clip", _build_clip_loader_cache_key), + "DualCLIPLoader": ("load_clip", _build_dual_clip_loader_cache_key), "VAELoader": ("load_vae", _build_vae_loader_cache_key), "unCLIPCheckpointLoader": ("load_checkpoint", _build_checkpoint_loader_cache_key), "ImageOnlyCheckpointLoader": ("load_checkpoint", _build_checkpoint_loader_cache_key), @@ -4194,10 +4217,18 @@ def _prepare_warm_container_for_request(volume: Any, payload: dict[str, Any]) -> 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, + ) return { "component_id": component_id, "task_id": os.getenv("MODAL_TASK_ID"), @@ -4210,6 +4241,85 @@ def _prepare_warm_container_for_request(volume: Any, payload: dict[str, Any]) -> } +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) + executed_plan_count = 0 + 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) + 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, + ) + executed_plan_count += 1 + except Exception: + if plan_key is not None: + with _LOADER_PREWARM_PLAN_KEYS_LOCK: + _LOADER_PREWARM_PLAN_KEYS.discard(plan_key) + raise + 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, + ) + + if modal is not None: # pragma: no branch - remote entrypoint configuration. settings = get_settings() app = modal.App(settings.app_name) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index ed2b648..6d80b30 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -5,6 +5,7 @@ - Prompt-scoped Modal class parameters are the wrong place to carry remote session identity. The `remote_session.session_id` belongs in the payload for data isolation, but `@app.cls` affinity should use reusable worker-pool slots like `worker-pool:slot:N`; otherwise every prompt asks Modal for a brand-new parameterized class instance and cross-prompt warm-container reuse never happens. - The local `custom_nodes/` tree is effectively process-static for one ComfyUI lifetime. Since adding or removing custom nodes still requires a ComfyUI restart before node availability changes, the queue-time sync engine should resolve `custom_nodes/` once per process and then reuse that cached decision instead of rehashing or rescanning the tree on every workflow run. - Warm bridge restoration is only half of the mapped seed optimization. Once every exported static bridge output has already been restored into the target remote session without replay, the seed payload should skip `prompt_executor_execute` entirely and synthesize its bridge-ref outputs directly from session state; otherwise the system still wastes seconds rerunning loader and encoder nodes after the expensive work was already avoided. +- Modal function snapshots are initialization snapshots, not mid-request checkpoints. The practical way to hide first-request model loads is to detect root literal loader nodes in the rewritten prompt, then execute them as synthetic one-node subgraphs during proactive warmup so the worker-local loader cache is hot before the real payload arrives. ## 2026-04-21 diff --git a/remote/modal_app.py b/remote/modal_app.py index 9bc8b15..a744a24 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -65,6 +65,14 @@ _PROMPT_WARMUP_STATES: dict[str, "_PromptWarmupState"] = {} _PROMPT_WARMUP_STATE_ORDER: queue.SimpleQueue[str] | None = None _PROMPT_WARMUP_STATE_CACHE_LIMIT = 256 _PRIMITIVE_WIDGET_INPUT_TYPES = frozenset({"INT", "FLOAT", "BOOLEAN", "STRING"}) +_ROOT_LOADER_PREWARM_CLASS_TYPES = frozenset( + { + "CheckpointLoaderSimple", + "UNETLoader", + "CLIPLoader", + "DualCLIPLoader", + } +) _BOUNDARY_INPUT_SIGNATURES_KEY = "__comfy_modal_boundary_input_signatures__" _REMOTE_CONTAINER_LOG_STREAMS_LOCK = threading.Lock() _REMOTE_CONTAINER_LOG_STREAMS: dict[str, "_RemoteContainerLogStreamState"] = {} @@ -3859,6 +3867,84 @@ def _clamp_prompt_warmup_target(warmup_target: int) -> int: 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 _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 _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any]: """Extract the prompt-scoped warmup-relevant fields from one payload.""" return { @@ -3874,6 +3960,7 @@ def _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any]: "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": _build_loader_prewarm_plans(payload), } diff --git a/settings.py b/settings.py index 27d8eb4..9775493 100644 --- a/settings.py +++ b/settings.py @@ -42,6 +42,7 @@ class ModalSyncSettings: max_containers: int | None = None buffer_containers: int | None = None enable_proactive_warmup: bool = True + enable_loader_prewarm: bool = True proactive_warmup_head_start_seconds: float = 2.0 stream_remote_container_logs: bool = False @@ -212,6 +213,7 @@ def get_settings() -> ModalSyncSettings: max_containers=_read_optional_int_env("COMFY_MODAL_MAX_CONTAINERS"), buffer_containers=_read_optional_int_env("COMFY_MODAL_BUFFER_CONTAINERS"), enable_proactive_warmup=_read_bool_env("COMFY_MODAL_ENABLE_PROACTIVE_WARMUP") is not False, + enable_loader_prewarm=_read_bool_env("COMFY_MODAL_ENABLE_LOADER_PREWARM") is not False, proactive_warmup_head_start_seconds=_read_float_env( "COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS", 2.0, diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index aad61d6..1ffa77e 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -703,6 +703,51 @@ def test_await_prompt_warmup_slots_waits_for_inflight_futures( assert completed_count == 1 +def test_build_prompt_warmup_request_includes_root_loader_prewarm_plans( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """Warmup requests should synthesize one-node plans for root literal loader nodes only.""" + monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true") + remote_modal_app_module.get_settings.cache_clear() + try: + warmup_request = remote_modal_app_module._build_prompt_warmup_request( + { + "prompt_id": "prompt-1", + "component_id": "component-1", + "subgraph_prompt": { + "1": { + "class_type": "UNETLoader", + "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"}, + }, + "2": { + "class_type": "CLIPLoader", + "inputs": {"clip_name": "clip-a.safetensors", "type": "stable_diffusion"}, + }, + "3": { + "class_type": "DualCLIPLoader", + "inputs": {"clip_name1": "clip-a.safetensors", "clip_name2": "clip-b.safetensors", "type": "flux"}, + }, + "4": { + "class_type": "UNETLoader", + "inputs": {"unet_name": ["99", 0]}, + }, + "5": { + "class_type": "KSampler", + "inputs": {"model": ["1", 0]}, + }, + }, + } + ) + finally: + remote_modal_app_module.get_settings.cache_clear() + + loader_plans = warmup_request["loader_prewarm_plans"] + assert [plan["node_id"] for plan in loader_plans] == ["1", "2", "3"] + assert [plan["class_type"] for plan in loader_plans] == ["UNETLoader", "CLIPLoader", "DualCLIPLoader"] + assert all(plan["execute_node_ids"] == [plan["node_id"]] for plan in loader_plans) + + def test_register_exact_component_parallelism_refines_prompt_target( remote_modal_app_module: Any, monkeypatch: Any, @@ -1890,6 +1935,7 @@ def test_modal_cloud_installs_loader_cache_wrappers_for_builtin_loaders( "CheckpointLoaderSimple": type("CheckpointLoaderSimple", (), {"load_checkpoint": lambda self, ckpt_name: (ckpt_name,)}), "UNETLoader": type("UNETLoader", (), {"load_unet": lambda self, unet_name, weight_dtype="default": (unet_name,)}), "CLIPLoader": type("CLIPLoader", (), {"load_clip": lambda self, clip_name, type="stable_diffusion", device="default": (clip_name,)}), + "DualCLIPLoader": type("DualCLIPLoader", (), {"load_clip": lambda self, clip_name1, clip_name2, type, device="default": (clip_name1, clip_name2)}), "VAELoader": type("VAELoader", (), {"load_vae": lambda self, vae_name: (vae_name,)}), "unCLIPCheckpointLoader": type("unCLIPCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}), "ImageOnlyCheckpointLoader": type("ImageOnlyCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}), @@ -1911,6 +1957,7 @@ def test_modal_cloud_installs_loader_cache_wrappers_for_builtin_loaders( "CheckpointLoaderSimple", "UNETLoader", "CLIPLoader", + "DualCLIPLoader", "VAELoader", "unCLIPCheckpointLoader", "ImageOnlyCheckpointLoader", @@ -2368,6 +2415,13 @@ def test_modal_cloud_prepares_warm_container_for_request( "_register_custom_nodes_root", lambda custom_nodes_root: calls.append(("register", custom_nodes_root)), ) + monkeypatch.setattr( + modal_cloud_module, + "_execute_loader_prewarm_plans", + lambda *, component_id, loader_prewarm_plans, custom_nodes_root: calls.append( + ("prewarm", component_id, tuple(plan["signature"] for plan in loader_prewarm_plans), custom_nodes_root) + ), + ) monkeypatch.setenv("MODAL_TASK_ID", "task-123") result = modal_cloud_module._prepare_warm_container_for_request( @@ -2378,12 +2432,14 @@ def test_modal_cloud_prepares_warm_container_for_request( "uploaded_volume_paths": ["/storage/example.bin"], "custom_nodes_bundle": "custom_nodes_bundle.zip", "warmup_slot_index": 2, + "loader_prewarm_plans": [{"signature": "loader-plan-1"}], }, ) assert calls == [ ("reload", "component-1", "marker-1", ["/storage/example.bin"]), ("register", Path("/tmp/extracted-bundle")), + ("prewarm", "component-1", ("loader-plan-1",), Path("/tmp/extracted-bundle")), ] assert result == { "component_id": "component-1", @@ -2393,6 +2449,61 @@ def test_modal_cloud_prepares_warm_container_for_request( } +def test_modal_cloud_executes_loader_prewarm_plans_once_per_worker( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """Worker-local loader prewarm plans should execute once and then be skipped on reuse.""" + original_plan_keys = set(modal_cloud_module._LOADER_PREWARM_PLAN_KEYS) + modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.clear() + observed_calls: list[tuple[str, tuple[str, ...]]] = [] + + monkeypatch.setattr( + modal_cloud_module, + "_ensure_comfy_runtime_initialized", + lambda custom_nodes_root: observed_calls.append(("runtime", (str(custom_nodes_root),))), + ) + monkeypatch.setattr( + modal_cloud_module, + "_execute_subgraph_prompt", + lambda payload, hydrated_inputs, custom_nodes_root, **kwargs: observed_calls.append( + ("execute", (str(payload["component_id"]), str(tuple(payload["execute_node_ids"])))) + ) or tuple(), + ) + monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true") + modal_cloud_module.get_settings.cache_clear() + try: + plans = [ + { + "signature": "loader-plan-1", + "node_id": "7", + "prompt_id": "prompt-1", + "subgraph_prompt": {"7": {"class_type": "UNETLoader", "inputs": {"unet_name": "model.safetensors"}}}, + "execute_node_ids": ["7"], + } + ] + modal_cloud_module._execute_loader_prewarm_plans( + component_id="component-1", + loader_prewarm_plans=plans, + custom_nodes_root=Path("/tmp/extracted-bundle"), + ) + modal_cloud_module._execute_loader_prewarm_plans( + component_id="component-1", + loader_prewarm_plans=plans, + custom_nodes_root=Path("/tmp/extracted-bundle"), + ) + finally: + modal_cloud_module.get_settings.cache_clear() + modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.clear() + modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.update(original_plan_keys) + + assert observed_calls == [ + ("runtime", ("/tmp/extracted-bundle",)), + ("execute", ("component-1::loader-prewarm:7", "('7',)")), + ("runtime", ("/tmp/extracted-bundle",)), + ] + + def test_remote_modal_auto_deploys_missing_app_by_default( remote_modal_app_module: Any, monkeypatch: Any,