fix: stamp snapshot profiles during prompt rewrite

README.md

@@ -174,6 +174,8 @@ Remote runtime behavior:
 - That snapshot profile now propagates through split static/mapped phase payloads as well, so mapped worker lanes can reuse the same snapshot-backed model set instead of silently falling back to the default no-profile worker class.
 - GPU snapshot-mode behavior is now passed into the remote `RemoteEngine` as an explicit class parameter, so `@modal.enter(snap=True)` no longer depends on the Modal container inheriting the local ComfyUI process environment variable.
 - Coarse payloads now persist any synthesized loader `snapshot_profile_key` back onto themselves before later split static/mapped payloads are derived, so mapped worker lanes inherit the same snapshot-backed model profile instead of silently recomputing from a loader-free per-item subgraph.
+- 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.
 - 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.

api_intercept.py

@@ -10,7 +10,7 @@ import time
 import uuid
 from collections import defaultdict, deque
 from dataclasses import dataclass, field
-from typing import Any, Iterator
+from typing import Any, Iterator, Mapping
 
 from aiohttp import web
 
@@ -41,6 +41,14 @@ _TRANSPORTABLE_OUTPUT_TYPES = frozenset(
         "STRING",
     }
 )
+_ROOT_LOADER_PREWARM_CLASS_TYPES = frozenset(
+    {
+        "CheckpointLoaderSimple",
+        "UNETLoader",
+        "CLIPLoader",
+        "DualCLIPLoader",
+    }
+)
 
 
 @dataclass(frozen=True)
@@ -365,6 +373,111 @@ def _prompt_node_signature_digest(
     return digest
 
 
+def _iter_loader_snapshot_prompt_payloads(payload: Mapping[str, Any]) -> Iterator[Mapping[str, Any]]:
+    """Yield prompt-bearing 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, Mapping):
+                yield phase_payload
+        return
+    if isinstance(split_proxy_payloads, list):
+        for phase_payload in split_proxy_payloads:
+            if isinstance(phase_payload, Mapping):
+                yield phase_payload
+        return
+    if isinstance(payload.get("subgraph_prompt"), Mapping):
+        yield payload
+
+
+def _is_root_literal_loader_node(prompt_node: Mapping[str, Any]) -> bool:
+    """Return whether one prompt node is a supported root loader with 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 _payload_loader_snapshot_profile_key(payload: Mapping[str, Any]) -> str:
+    """Return the stable loader snapshot profile key derivable from one payload."""
+    prompt_id = payload.get("prompt_id")
+    normalized_prompt_id = str(prompt_id) if prompt_id is not None else None
+    plan_signatures: set[str] = set()
+    for prompt_payload in _iter_loader_snapshot_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
+            plan_signatures.add(_loader_prewarm_plan_signature(class_type, inputs))
+            logger.debug(
+                "Derived rewrite-time loader prewarm plan for component=%s node=%s class_type=%s prompt_id=%s.",
+                payload.get("component_id"),
+                node_id,
+                class_type,
+                normalized_prompt_id,
+            )
+    if not plan_signatures:
+        return ""
+    profile_digest = hashlib.sha256(
+        json.dumps({"plan_signatures": sorted(plan_signatures)}, sort_keys=True).encode("utf-8")
+    ).hexdigest()
+    return f"loader-profile:{profile_digest}"
+
+
+def _stamp_snapshot_profile_key(payload: dict[str, Any], snapshot_profile_key: str) -> None:
+    """Attach one loader snapshot profile key to a payload and any split descendants."""
+    if not snapshot_profile_key:
+        return
+    payload["snapshot_profile_key"] = snapshot_profile_key
+    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):
+                phase_payload["snapshot_profile_key"] = snapshot_profile_key
+        return
+    if isinstance(split_proxy_payloads, list):
+        for phase_payload in split_proxy_payloads:
+            if isinstance(phase_payload, dict):
+                phase_payload["snapshot_profile_key"] = snapshot_profile_key
+
+
+def _attach_snapshot_profile_key(payload: dict[str, Any], settings: ModalSyncSettings) -> dict[str, Any]:
+    """Stamp a deterministic loader snapshot profile onto one payload when enabled."""
+    if not settings.enable_gpu_memory_snapshot or not settings.enable_loader_prewarm:
+        return payload
+    snapshot_profile_key = _payload_loader_snapshot_profile_key(payload)
+    if snapshot_profile_key:
+        _stamp_snapshot_profile_key(payload, snapshot_profile_key)
+        logger.info(
+            "Attached rewrite-time loader snapshot profile %s to component=%s payload_kind=%s.",
+            snapshot_profile_key,
+            payload.get("component_id"),
+            payload.get("payload_kind"),
+        )
+    return payload
+
+
 def _boundary_source_signature(
     prompt: dict[str, Any],
     source: LinkedOutputRef,
@@ -1812,6 +1925,7 @@ def _build_component_payload(
     component_prompt: dict[str, Any],
     signature_prompt: dict[str, Any],
     extra_data: dict[str, Any] | None,
+    settings: ModalSyncSettings,
     requires_volume_reload: bool,
     volume_reload_marker: str | None,
     custom_nodes_bundle: SyncedAsset | None,
@@ -2072,7 +2186,10 @@ def _build_component_payload(
 
     split_phase_payloads = build_phase_payloads_for_transportable_splits()
     if split_phase_payloads is not None:
-        return {"split_proxy_payloads": split_phase_payloads}
+        return _attach_snapshot_profile_key(
+            {"split_proxy_payloads": split_phase_payloads},
+            settings,
+        )
 
     payload = {
         "payload_kind": "mapped_subgraph" if component.mapped_boundary_input_name else "subgraph",
@@ -2235,7 +2352,7 @@ def _build_component_payload(
             "execute_node_ids": list(component.mapped_execute_node_ids),
         }
         if not component.static_node_ids:
-            return payload
+            return _attach_snapshot_profile_key(payload, settings)
         remote_session = RemoteSessionHandle(
             session_id=uuid.uuid4().hex,
             prompt_id=(str(prompt_id) if prompt_id is not None else None),
@@ -2330,7 +2447,7 @@ def _build_component_payload(
                 )
             ),
         }
-    return payload
+    return _attach_snapshot_profile_key(payload, settings)
 
 
 def _component_uploaded_volume_paths(
@@ -2739,6 +2856,7 @@ def rewrite_prompt_for_modal(
             component_prompt=synced_component_prompts[component.representative_node_id],
             signature_prompt=prompt,
             extra_data=extra_data,
+            settings=resolved_settings,
             requires_volume_reload=bool(uploaded_volume_paths),
             volume_reload_marker=volume_reload_marker,
             custom_nodes_bundle=summary.custom_nodes_bundle,

docs/LEARNINGS.md

@@ -152,6 +152,8 @@
 - If the goal is faster cold restarts for heavyweight diffusion loaders, request-time warmup is not enough. The loader work has to move into `@modal.enter(snap=True)` behind a stable per-model-set snapshot profile, otherwise fresh workers will still pay one `UNETLoader`/`CLIPLoader` miss each even when GPU snapshots are enabled.
 - That snapshot profile must also survive any split-payload rewriting. If static/mapped phase payload builders drop `snapshot_profile_key`, mapped worker lanes will look up the default `RemoteEngine` class and all the snapshot work will appear to do nothing in logs.
 - Preserving `snapshot_profile_key` in split-phase helpers is still not enough if the key only exists as a transient local variable during deployed-class lookup. Once a coarse payload synthesizes a loader snapshot profile, that key has to be written back onto the payload itself so later mapped/static descendants inherit it instead of recomputing from a loader-free subgraph and falling back to `snapshot_profile=None`.
+- 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.
 - 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.

remote/modal_app.py

@@ -3709,12 +3709,18 @@ def _lookup_deployed_remote_engine(
     affinity_key = affinity_key_override
     if affinity_key is None:
         affinity_key = _remote_worker_affinity_key(payload)
+    loader_prewarm_plans: list[dict[str, Any]] = _build_loader_prewarm_plans(payload)
     snapshot_profile_key = ""
     payload_snapshot_profile_key = payload.get("snapshot_profile_key")
     if isinstance(payload_snapshot_profile_key, str):
         snapshot_profile_key = payload_snapshot_profile_key.strip()
-    if not snapshot_profile_key:
-        snapshot_profile_key = _store_loader_snapshot_profile(_build_loader_prewarm_plans(payload))
+    if snapshot_profile_key:
+        stored_snapshot_profile_key = _store_loader_snapshot_profile(loader_prewarm_plans)
+        if stored_snapshot_profile_key:
+            snapshot_profile_key = stored_snapshot_profile_key
+            payload["snapshot_profile_key"] = snapshot_profile_key
+    else:
+        snapshot_profile_key = _store_loader_snapshot_profile(loader_prewarm_plans)
         if snapshot_profile_key:
             payload["snapshot_profile_key"] = snapshot_profile_key
     logger.info(

tests/test_api_intercept.py

@@ -16,6 +16,14 @@ class _FakeRemoteModelNode:
     OUTPUT_IS_LIST = (False,)
 
 
+class _FakeCheckpointLoaderSimpleNode:
+    """Fake root loader node that produces a non-transportable MODEL output."""
+
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("model",)
+    OUTPUT_IS_LIST = (False,)
+
+
 class _FakeRemoteSamplerNode:
     """Fake node that consumes a model and produces a transportable latent."""
 
@@ -1390,6 +1398,116 @@ def test_rewrite_supports_mapped_branch_that_shares_non_transportable_upstream_w
     assert rewritten_prompt["8"]["inputs"]["image"] == ["1__mapped", 0]
 
 
+def test_rewrite_stamps_snapshot_profile_on_split_static_and_mapped_payloads(
+    api_intercept_module: Any,
+    settings_module: Any,
+    sync_engine_module: Any,
+    tmp_path: Path,
+) -> None:
+    """Split static and mapped payloads should inherit the same loader snapshot profile."""
+    settings = settings_module.ModalSyncSettings(
+        app_name="app",
+        auto_deploy=True,
+        allow_ephemeral_fallback=False,
+        enable_memory_snapshot=True,
+        enable_gpu_memory_snapshot=True,
+        execution_mode="local",
+        sync_custom_nodes=False,
+        volume_name="volume",
+        route_path="/modal/queue_prompt",
+        marker_property="is_modal_remote",
+        local_storage_root=tmp_path / "storage",
+        remote_storage_root="/storage",
+        custom_nodes_archive_name="custom_nodes_bundle.zip",
+        comfyui_root=None,
+        custom_nodes_dir=tmp_path / "custom_nodes",
+    )
+    sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
+    fake_nodes_module = type(
+        "FakeNodesModule",
+        (),
+        {
+            "NODE_CLASS_MAPPINGS": {
+                "CheckpointLoaderSimple": _FakeCheckpointLoaderSimpleNode,
+                "RemoteSampler": _FakeRemoteSamplerNode,
+                "LatentSource": _FakeLatentSourceNode,
+                "ModalMapInput": _FakeModalMapInputNode,
+                "LocalSink": _FakeLocalSinkNode,
+            },
+            "NODE_DISPLAY_NAME_MAPPINGS": {},
+        },
+    )()
+
+    workflow = {
+        "nodes": [
+            {"id": 1, "properties": {"is_modal_remote": True}},
+            {"id": 2, "properties": {"is_modal_remote": False}},
+            {"id": 3, "properties": {"is_modal_remote": True}},
+            {"id": 4, "properties": {"is_modal_remote": False}},
+            {"id": 5, "properties": {"is_modal_remote": False}},
+            {"id": 6, "properties": {"is_modal_remote": True}},
+            {"id": 7, "properties": {"is_modal_remote": True}},
+            {"id": 8, "properties": {"is_modal_remote": False}},
+        ]
+    }
+    prompt = {
+        "1": {
+            "class_type": "CheckpointLoaderSimple",
+            "inputs": {"ckpt_name": "base.safetensors"},
+            "_meta": {"title": "Checkpoint"},
+        },
+        "2": {
+            "class_type": "LatentSource",
+            "inputs": {},
+            "_meta": {"title": "Single Latent"},
+        },
+        "3": {
+            "class_type": "RemoteSampler",
+            "inputs": {"model": ["1", 0], "latent": ["2", 0]},
+            "_meta": {"title": "Unmapped Sampler"},
+        },
+        "4": {
+            "class_type": "LocalSink",
+            "inputs": {"image": ["3", 0]},
+            "_meta": {"title": "Local Sink 1"},
+        },
+        "5": {
+            "class_type": "LatentSource",
+            "inputs": {},
+            "_meta": {"title": "Batch Latent Source"},
+        },
+        "6": {
+            "class_type": "ModalMapInput",
+            "inputs": {"value": ["5", 0]},
+            "_meta": {"title": "Map Input"},
+        },
+        "7": {
+            "class_type": "RemoteSampler",
+            "inputs": {"model": ["1", 0], "latent": ["6", 0]},
+            "_meta": {"title": "Mapped Sampler"},
+        },
+        "8": {
+            "class_type": "LocalSink",
+            "inputs": {"image": ["7", 0]},
+            "_meta": {"title": "Local Sink 2"},
+        },
+    }
+
+    rewritten_prompt, _ = api_intercept_module.rewrite_prompt_for_modal(
+        prompt=prompt,
+        workflow=workflow,
+        sync_engine=sync_engine,
+        settings=settings,
+        nodes_module=fake_nodes_module,
+    )
+
+    static_payload = rewritten_prompt["1"]["inputs"]["original_node_data"]
+    mapped_payload = rewritten_prompt["1__mapped"]["inputs"]["original_node_data"]
+
+    assert static_payload["snapshot_profile_key"].startswith("loader-profile:")
+    assert mapped_payload["snapshot_profile_key"] == static_payload["snapshot_profile_key"]
+
+
 def test_rewrite_splits_unmapped_remote_siblings_that_share_non_transportable_upstream(
     api_intercept_module: Any,
     modal_executor_module: Any,

tests/test_modal_executor_node.py

@@ -3038,6 +3038,78 @@ def test_lookup_deployed_remote_engine_passes_snapshot_profile_parameter_for_gpu
     assert payload["snapshot_profile_key"] == result["snapshot_profile_key"]
 
 
+def test_lookup_deployed_remote_engine_stores_existing_snapshot_profile_record_when_possible(
+    remote_modal_app_module: Any,
+    monkeypatch: Any,
+) -> None:
+    """Deployed lookups should backfill snapshot records even when the payload already carries the key."""
+    observed_kwargs: list[dict[str, Any]] = []
+    snapshot_profiles: dict[str, Any] = {}
+
+    class FakeModal:
+        """Minimal modal SDK double that captures class construction kwargs."""
+
+        class Dict:
+            """Namespace for fake dict lookups."""
+
+            @staticmethod
+            def from_name(dict_name: str, create_if_missing: bool = False) -> Any:
+                """Return the shared fake snapshot profile store."""
+                assert dict_name == "comfy-modal-sync-snapshot-profiles"
+                assert create_if_missing is True
+                return snapshot_profiles
+
+        class Cls:
+            """Namespace for deployed class lookups."""
+
+            @staticmethod
+            def from_name(app_name: str, class_name: str) -> Any:
+                """Return a factory that records the provided class parameters."""
+                assert app_name == "comfy-modal-sync"
+                assert class_name == "RemoteEngine"
+
+                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
+                    """Record one synthesized Modal class constructor call."""
+                    observed_kwargs.append(kwargs)
+                    return kwargs
+
+                return build_remote_engine
+
+    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
+    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
+    remote_modal_app_module.get_settings.cache_clear()
+    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
+    expected_snapshot_profile_key = remote_modal_app_module._loader_snapshot_profile_key(
+        [
+            {
+                "signature": remote_modal_app_module._loader_prewarm_plan_signature(
+                    "UNETLoader",
+                    {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
+                )
+            }
+        ]
+    )
+    payload = {
+        "component_id": "component-1",
+        "snapshot_profile_key": expected_snapshot_profile_key,
+        "subgraph_prompt": {
+            "1": {
+                "class_type": "UNETLoader",
+                "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
+            }
+        },
+    }
+    try:
+        result = remote_modal_app_module._lookup_deployed_remote_engine(payload)
+    finally:
+        remote_modal_app_module.get_settings.cache_clear()
+        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
+
+    assert result["snapshot_profile_key"] == expected_snapshot_profile_key
+    assert observed_kwargs == [result]
+    assert expected_snapshot_profile_key in snapshot_profiles
+
+
 def test_lookup_deployed_remote_engine_reuses_worker_pool_slots_across_prompt_sessions(
     remote_modal_app_module: Any,
     monkeypatch: Any,