perf: prewarm root loader nodes during modal warmup

This commit is contained in:
Ken Simpson
2026-04-22 15:41:49 -07:00
parent 9997f56985
commit e2276dae19
6 changed files with 313 additions and 0 deletions
+2
View File
@@ -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
+110
View File
@@ -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)
+1
View File
@@ -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
+87
View File
@@ -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),
}
+2
View File
@@ -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,
+111
View File
@@ -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,