diff --git a/README.md b/README.md index 0f336c5..85822f8 100644 --- a/README.md +++ b/README.md @@ -177,6 +177,7 @@ Remote runtime behavior: - The queue rewrite now stamps that same loader `snapshot_profile_key` onto already-materialized split proxy payloads up front, because runtime lookup happens too late to fix payloads that were cloned earlier during graph rewrite. - Remote lookup now also backfills the snapshot-profile record when a payload already carries the key but still has enough loader context to reconstruct the record, which keeps rewrite-time stamped profiles usable on first lookup. - Generic `snapshot_profile=None` workers now skip GPU snapshot runtime prewarm entirely. Only profiled `loader-profile:...` workers do full ComfyUI initialization in `snap=True`, which avoids wasting time on unstable no-profile GPU snapshots that do not help cold-start model loading anyway. +- When GPU snapshots are enabled, proactive `::warmup` lookups are now skipped entirely for no-profile payloads. Modal only warms concrete profiled workers, so the scheduler no longer burns time on generic `snapshot_profile=None` classes that cannot produce useful model-loaded snapshots. - The default `scaledown_window` is `600` seconds with `min_containers=0`. - Warm containers can reuse loader state and `PromptExecutor` state across compatible requests. - Each Modal GPU container now handles one active workflow execution at a time. If multiple remote components become ready in parallel, Modal can scale them out across multiple containers instead of multiplexing several executions onto one GPU worker. @@ -394,7 +395,7 @@ If a remote-marked node depends on a model filename that cannot be resolved to a ### Modal runtime sizing - `COMFY_MODAL_ENABLE_MEMORY_SNAPSHOT`: Enable Modal CPU memory snapshots. Default: `true`. -- `COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT`: Enable Modal GPU memory snapshots. Default: `false`. +- `COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT`: Enable Modal GPU memory snapshots. Default: `true`. - `COMFY_MODAL_SNAPSHOT_PROFILE_DICT_NAME`: Modal `Dict` used to publish per-model-set snapshot loader plans to `snap=True`. Default: `-snapshot-profiles`. - `COMFY_MODAL_GPU`: Modal GPU type to request when deploying the remote class. Default: `A100`. - `COMFY_MODAL_SCALEDOWN_WINDOW`: Keep idle containers warm for this many seconds. Default: `600`. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 0572a92..6034ac7 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -155,6 +155,7 @@ - Even that runtime-side coarse-payload writeback can still be too late when the queue rewrite has already materialized split proxy payloads. In that case the rewrite step itself must stamp the derived `snapshot_profile_key` onto the split payloads before they are inserted into the rewritten prompt. - A payload carrying `snapshot_profile_key` is not sufficient by itself if the underlying snapshot-profile record has never been stored. Remote lookup should opportunistically store the record whenever it can still reconstruct the loader prewarm plan set, even if the payload already arrived with the key prefilled. - GPU snapshots are not worth enabling blindly for generic no-profile workers. A `snapshot_profile=None` container that snapshots only “ComfyUI core initialized” state can still be unstable after restore and does not buy the model-loaded cold-start win we care about, so the safer policy is to reserve GPU snapshot prewarm for real `loader-profile:...` workers only. +- The same rule should apply to proactive warmup. If GPU snapshots are enabled but `_build_prompt_warmup_request()` cannot derive a loader profile, the scheduler should skip the `::warmup` lookup entirely instead of instantiating a generic `snapshot_profile=None` class that cannot create a useful model-loaded snapshot. - Deployed Modal class behavior that matters at `snap=True` should not be inferred from environment inside the remote container. The container does not automatically inherit the local ComfyUI launch env, so flags like GPU-snapshot mode need to be carried as deploy-time/class parameters if the snapshot hook logic depends on them. - A remote proxy cannot treat every incoming list literal as just another scalar boundary input. Ordinary remote components also have to preserve ComfyUI's implicit zipped batch semantics across the boundary, or primitive widget sockets like `seed: INT` end up receiving whole lists such as `[10, 11, 12, 13]` instead of per-item values. - That implicit batching path also needs the same hybrid static-vs-mapped split as explicit `ModalMapInput`. If two remote execute nodes share a non-transportable upstream object like `MODEL` but only one of them depends on the multi-item boundary input, the sibling execute node still has to run once, not once per batch item. diff --git a/remote/modal_app.py b/remote/modal_app.py index 6d15214..e235027 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -4029,9 +4029,17 @@ def _build_loader_prewarm_plans(payload: dict[str, Any]) -> list[dict[str, Any]] return plans -def _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any]: +def _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any] | None: """Extract the prompt-scoped warmup-relevant fields from one payload.""" + settings = get_settings() loader_prewarm_plans = _build_loader_prewarm_plans(payload) + snapshot_profile_key = _store_loader_snapshot_profile(loader_prewarm_plans) + if settings.enable_gpu_memory_snapshot and not snapshot_profile_key: + 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")) @@ -4046,7 +4054,7 @@ def _build_prompt_warmup_request(payload: dict[str, Any]) -> dict[str, Any]: "uploaded_volume_paths": list(payload.get("uploaded_volume_paths", [])), "custom_nodes_bundle": payload.get("custom_nodes_bundle"), "loader_prewarm_plans": loader_prewarm_plans, - "snapshot_profile_key": _store_loader_snapshot_profile(loader_prewarm_plans), + "snapshot_profile_key": snapshot_profile_key, } @@ -4283,7 +4291,7 @@ async def _await_prompt_warmup_slots( def ensure_remote_warm_capacity( - warmup_request: dict[str, Any], + warmup_request: dict[str, Any] | None, *, warmup_target: int, reason: str, @@ -4294,6 +4302,8 @@ def ensure_remote_warm_capacity( 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: diff --git a/settings.py b/settings.py index b50796d..7de751e 100644 --- a/settings.py +++ b/settings.py @@ -172,7 +172,8 @@ def get_settings() -> ModalSyncSettings: auto_deploy=_read_bool_env("COMFY_MODAL_AUTO_DEPLOY") is not False, allow_ephemeral_fallback=_read_bool_env("COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK") or False, enable_memory_snapshot=_read_bool_env("COMFY_MODAL_ENABLE_MEMORY_SNAPSHOT") is not False, - enable_gpu_memory_snapshot=_read_bool_env("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT") or False, + enable_gpu_memory_snapshot=_read_bool_env("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT") + is not False, execution_mode=execution_mode, sync_custom_nodes=sync_custom_nodes, volume_name=os.getenv("COMFY_MODAL_VOLUME_NAME", "comfy-universal-storage"), diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 5b8e88d..2e9b258 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -709,6 +709,7 @@ def test_build_prompt_warmup_request_includes_root_loader_prewarm_plans( ) -> None: """Warmup requests should synthesize one-node plans for root literal loader nodes only.""" monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true") + monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "false") remote_modal_app_module.get_settings.cache_clear() try: warmup_request = remote_modal_app_module._build_prompt_warmup_request( @@ -804,6 +805,33 @@ def test_build_prompt_warmup_request_registers_snapshot_profile_when_gpu_snapsho assert snapshot_profiles[snapshot_profile_key]["loader_prewarm_plans"] == warmup_request["loader_prewarm_plans"] +def test_build_prompt_warmup_request_skips_generic_gpu_snapshot_warmup_without_profile( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """Generic proactive warmup should be skipped when GPU snapshots are enabled without a loader profile.""" + monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true") + 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": { + "5": { + "class_type": "KSampler", + "inputs": {"model": ["1", 0]}, + } + }, + } + ) + finally: + remote_modal_app_module.get_settings.cache_clear() + + assert warmup_request is None + + def test_register_exact_component_parallelism_refines_prompt_target( remote_modal_app_module: Any, monkeypatch: Any, @@ -2964,12 +2992,12 @@ def test_lookup_deployed_remote_engine_passes_affinity_as_modal_parameter( assert result == { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, } assert observed_kwargs == [ { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, } ] @@ -3163,20 +3191,20 @@ def test_lookup_deployed_remote_engine_reuses_worker_pool_slots_across_prompt_se assert first_result == { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, } assert second_result == { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, } assert observed_kwargs == [ { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, }, { "worker_affinity_key": "worker-pool:slot:0", - "gpu_snapshot_enabled": False, + "gpu_snapshot_enabled": True, }, ] diff --git a/tests/test_settings.py b/tests/test_settings.py index 3b8a8fe..5060da4 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -59,6 +59,21 @@ def test_settings_reads_modal_gpu_override( assert settings.modal_gpu == "L40S" +def test_settings_enable_gpu_memory_snapshot_defaults_true( + settings_module: Any, + monkeypatch: Any, +) -> None: + """GPU memory snapshots should be enabled by default.""" + monkeypatch.delenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", raising=False) + settings_module.get_settings.cache_clear() + try: + settings = settings_module.get_settings() + finally: + settings_module.get_settings.cache_clear() + + assert settings.enable_gpu_memory_snapshot is True + + def test_settings_reads_modal_container_scaling_overrides( settings_module: Any, monkeypatch: Any,