Reserve Modal cooldown for exactly one worker pool per prompt

Worker slots live in separate namespaces (mapped lanes use
worker-pool:slot:N, ordinary components worker-pool:comfy|llm:slot:N),
each its own Modal container. Matching only the trailing slot:N gave the
reserve cooldown to both the static comfy:slot:0 worker and mapped lane
slot:0 in hybrid prompts, so reserve_cooldown_count=1 kept two idle GPUs.

The planner now stamps one modal_reserve_cooldown_pool per Modal
environment (mapped lanes when the prompt has mapped work, else comfy,
else llm), and the lookup grants the reserve cooldown only to leading
slots of that pool; every other worker uses the global cooldown.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

README.md

@@ -143,7 +143,7 @@ The Modal Configuration node's advanced settings control how long idle workers s
 | Setting | Default | Meaning |
 | --- | --- | --- |
 | `global_cooldown` | `60` | Idle seconds before any Modal worker in this pool scales down. |
-| `reserve_cooldown_count` | `0` | How many workers (the first N worker slots, at most `instance_count`) use `reserve_cooldown` instead. |
+| `reserve_cooldown_count` | `0` | How many workers use `reserve_cooldown` instead (at most `instance_count`). They are the first N slots of one worker pool per prompt: the mapped lanes when the prompt has mapped work, otherwise the ordinary ComfyUI workers, otherwise the LLM workers. |
 | `reserve_cooldown` | `300` | Idle seconds before a reserved worker scales down. |
 
 Set a small `reserve_cooldown_count` with a long `reserve_cooldown` to keep a few always-warm workers while the rest of the pool releases its GPUs quickly. Cooldowns apply per call (Modal `with_options(scaledown_window=...)`), so changing them never redeploys the app. They override `COMFY_MODAL_SCALEDOWN_WINDOW` for workflows that use a Modal Configuration node; local-gap pools still stay warm for at least the local keepalive window.

docs/LEARNINGS.md

@@ -793,3 +793,16 @@ nested executor inside the host's running prompt, so it now passes zeroed
 classic-cache `cache_args`. It also restores the host's release callback and
 headroom afterwards; otherwise the nested run would silently switch off the
 host executor's RAM-pressure eviction for the rest of the outer prompt.
+
+### Reserve cooldown must target one worker pool
+
+Worker slots live in separate namespaces: mapped lanes use `worker-pool:slot:N`,
+while ordinary components use `worker-pool:comfy:slot:N` or `worker-pool:llm:slot:N`.
+Each is its own Modal container. The first cooldown version keyed only on the
+trailing `slot:N`, so in a hybrid static+mapped prompt both the static
+`comfy:slot:0` worker and mapped lane `slot:0` held the 300s reserve: two idle
+GPUs for `reserve_cooldown_count=1`, and five containers for `instance_count=4`.
+`_configure_reserve_cooldown_pools` now stamps one `modal_reserve_cooldown_pool`
+per Modal environment (mapped > comfy > llm), and the lookup gives the reserve
+cooldown only to slots in that pool.
+

prompt_affinity_planning.py

@@ -24,7 +24,14 @@ if __package__:
         _node_output_refs,
     )
     from .remote_plan_types import InputTarget, LinkedOutputRef
-    from .settings import MODAL_INHERITED_PAYLOAD_FIELDS
+    from .settings import (
+        MODAL_GLOBAL_COOLDOWN_FIELD,
+        MODAL_INHERITED_PAYLOAD_FIELDS,
+        MODAL_RESERVE_COOLDOWN_POOL_FIELD,
+        MODAL_WORKER_POOL_COMFY,
+        MODAL_WORKER_POOL_LLM,
+        MODAL_WORKER_POOL_MAPPED,
+    )
 else:  # pragma: no cover - flat import inside the Modal container.
     from component_planning import _component_ancestors_of_local_source
     from execution_environments import ExecutionProvider
@@ -41,7 +48,14 @@ else:  # pragma: no cover - flat import inside the Modal container.
         _node_output_refs,
     )
     from remote_plan_types import InputTarget, LinkedOutputRef
-    from settings import MODAL_INHERITED_PAYLOAD_FIELDS
+    from settings import (
+        MODAL_GLOBAL_COOLDOWN_FIELD,
+        MODAL_INHERITED_PAYLOAD_FIELDS,
+        MODAL_RESERVE_COOLDOWN_POOL_FIELD,
+        MODAL_WORKER_POOL_COMFY,
+        MODAL_WORKER_POOL_LLM,
+        MODAL_WORKER_POOL_MAPPED,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -385,6 +399,93 @@ def _speculative_prewarm_target_payload(
     }
 
 
+def _reserve_cooldown_pool(
+    component_ids: list[str],
+    payloads_by_component_id: Mapping[str, Mapping[str, Any]],
+    mapped_component_ids: set[str],
+) -> str:
+    """Pick the single worker pool whose leading slots keep the reserve cooldown.
+
+    Mapped lanes do the bulk of the work when present; otherwise ordinary ComfyUI
+    workers, and LLM workers only when nothing else runs on this Modal target.
+    """
+    if any(component_id in mapped_component_ids for component_id in component_ids):
+        return MODAL_WORKER_POOL_MAPPED
+    affinity_groups = {
+        str(
+            payloads_by_component_id[component_id].get("remote_worker_affinity_group")
+            or MODAL_WORKER_POOL_COMFY
+        )
+        .strip()
+        .lower()
+        for component_id in component_ids
+    }
+    if MODAL_WORKER_POOL_COMFY in affinity_groups:
+        return MODAL_WORKER_POOL_COMFY
+    return MODAL_WORKER_POOL_LLM
+
+
+def _cooldown_payloads_by_environment(
+    rewritten_prompt: dict[str, Any],
+    component_ids: list[str],
+) -> tuple[dict[str, Mapping[str, Any]], dict[str, list[str]]]:
+    """Return Modal payloads carrying a cooldown policy, grouped by environment."""
+    payloads_by_component_id: dict[str, Mapping[str, Any]] = {}
+    component_ids_by_environment: dict[str, list[str]] = {}
+    for component_id in component_ids:
+        payload = _remote_proxy_payload(rewritten_prompt, component_id)
+        if (
+            payload is None
+            or payload.get("execution_provider") != ExecutionProvider.MODAL.value
+            or payload.get(MODAL_GLOBAL_COOLDOWN_FIELD) is None
+        ):
+            continue
+        payloads_by_component_id[component_id] = payload
+        environment_id = str(payload.get("execution_environment_id") or "")
+        component_ids_by_environment.setdefault(environment_id, []).append(
+            component_id
+        )
+    return payloads_by_component_id, component_ids_by_environment
+
+
+def _configure_reserve_cooldown_pools(
+    *,
+    rewritten_prompt: dict[str, Any],
+    component_ids: list[str],
+    mapped_component_ids: set[str],
+) -> None:
+    """Stamp each Modal payload with the one pool that holds reserved workers.
+
+    Without this, slot 0 of the static ``comfy`` pool and slot 0 of the mapped
+    lanes would both keep the reserve cooldown, doubling the idle GPUs paid for.
+    """
+    payloads_by_component_id, component_ids_by_environment = (
+        _cooldown_payloads_by_environment(rewritten_prompt, component_ids)
+    )
+    for environment_id, environment_component_ids in (
+        component_ids_by_environment.items()
+    ):
+        reserve_pool = _reserve_cooldown_pool(
+            environment_component_ids,
+            payloads_by_component_id,
+            mapped_component_ids,
+        )
+        for component_id in environment_component_ids:
+            rewritten_prompt[component_id]["inputs"][
+                "original_node_data"
+            ] = update_registered_proxy_payload_fields(
+                component_id,
+                payloads_by_component_id[component_id],
+                {MODAL_RESERVE_COOLDOWN_POOL_FIELD: reserve_pool},
+            )
+        logger.info(
+            "Assigned Modal reserve cooldown pool=%s environment=%s components=%s.",
+            reserve_pool,
+            environment_id or "<default>",
+            environment_component_ids,
+        )
+
+
 def _configure_speculative_affinity_prewarm_payloads(
     *,
     rewritten_prompt: dict[str, Any],

prompt_interception.py

@@ -39,6 +39,7 @@ if __package__:
     )
     from .prompt_affinity_planning import (
         _configure_local_gap_keepalive_payloads,
+        _configure_reserve_cooldown_pools,
         _configure_speculative_affinity_prewarm_payloads,
         _parallelize_non_returning_local_branches,
     )
@@ -102,6 +103,7 @@ else:  # pragma: no cover - flat import inside the Modal container.
     )
     from prompt_affinity_planning import (
         _configure_local_gap_keepalive_payloads,
+        _configure_reserve_cooldown_pools,
         _configure_speculative_affinity_prewarm_payloads,
         _parallelize_non_returning_local_branches,
     )
@@ -974,6 +976,11 @@ def _finalize_rewritten_prompt(state: _PromptRewriteState) -> None:
     execution_stages = _component_execution_stages(
         state.rewritten_prompt, proxy_groups
     )
+    _configure_reserve_cooldown_pools(
+        rewritten_prompt=state.rewritten_prompt,
+        component_ids=list(state.summary.remote_component_ids),
+        mapped_component_ids=state.mapped_proxy_component_ids,
+    )
     _configure_speculative_affinity_prewarm_payloads(
         rewritten_prompt=state.rewritten_prompt,
         execution_stages=execution_stages,

remote/modal_deployment.py

@@ -27,6 +27,8 @@ from ..runtime_environment import (
     build_remote_runtime_identity,
 )
 from ..settings import (
+    MODAL_RESERVE_COOLDOWN_POOL_FIELD,
+    MODAL_WORKER_POOL_MAPPED,
     ModalCooldownPolicy,
     ModalSyncSettings,
     get_settings,
@@ -54,7 +56,9 @@ _MODAL_AUTO_DEPLOY_LOCK = threading.Lock()
 _MODAL_AUTO_DEPLOY_STATES: dict[tuple[str, str | None], "_ModalAutoDeployState"] = {}
 _MODAL_REMOTE_APP_VERSION_OK: set[tuple[str, str | None, str]] = set()
 _MODAL_APP_STOP_TIMEOUT_SECONDS = 120.0
-_AFFINITY_SLOT_PATTERN = re.compile(r"(?:^|:)slot:(\d+)$")
+_AFFINITY_SLOT_PATTERN = re.compile(
+    r"^worker-pool:(?:(?P<pool>[a-z]+):)?slot:(?P<slot>\d+)$"
+)
 
 
 @dataclass(frozen=True)
@@ -199,12 +203,15 @@ def _mapped_lane_affinity_key(payload: dict[str, Any], lane_index: int) -> str |
     del payload
     return _remote_worker_pool_affinity_key(lane_index)
 
-def _affinity_key_slot_index(worker_affinity_key: str) -> int | None:
-    """Return the worker-pool slot encoded in an affinity key, if it has one."""
-    match = _AFFINITY_SLOT_PATTERN.search(worker_affinity_key)
+def _affinity_key_worker(worker_affinity_key: str) -> tuple[str | None, int | None]:
+    """Return the (worker pool, slot) encoded in an affinity key, or (None, None).
+
+    Plain ``worker-pool:slot:N`` keys belong to the mapped-lane pool.
+    """
+    match = _AFFINITY_SLOT_PATTERN.fullmatch(worker_affinity_key)
     if match is None:
-        return None
-    return int(match.group(1))
+        return None, None
+    return match.group("pool") or MODAL_WORKER_POOL_MAPPED, int(match.group("slot"))
 
 def _scaledown_window_override(
     payload: Mapping[str, Any],
@@ -213,14 +220,21 @@ def _scaledown_window_override(
 ) -> int | None:
     """Return the per-call scaledown window, or None to keep the deployed default.
 
-    A workflow cooldown policy picks the reserve or global cooldown by worker
-    slot; local-gap pools are additionally held for the local keepalive window.
+    A workflow cooldown policy gives the reserve cooldown only to the leading
+    slots of the prompt's reserve pool; local-gap pools are additionally held for
+    the local keepalive window.
     """
     cooldown_policy = ModalCooldownPolicy.from_payload(payload)
     scaledown_window: int | None = None
     if cooldown_policy is not None:
-        scaledown_window = cooldown_policy.scaledown_window_for_slot(
-            _affinity_key_slot_index(worker_affinity_key)
+        worker_pool, slot_index = _affinity_key_worker(worker_affinity_key)
+        scaledown_window = cooldown_policy.scaledown_window_for_worker(
+            worker_pool=worker_pool,
+            slot_index=slot_index,
+            reserve_pool=str(
+                payload.get(MODAL_RESERVE_COOLDOWN_POOL_FIELD)
+                or MODAL_WORKER_POOL_MAPPED
+            ),
         )
     if bool(payload.get("remote_local_gap_pool")):
         scaledown_window = max(

settings.py

@@ -261,10 +261,12 @@ copy these, or it will expect a different fingerprint for the same Modal app.
 MODAL_GLOBAL_COOLDOWN_FIELD = "modal_global_cooldown_seconds"
 MODAL_RESERVE_COOLDOWN_COUNT_FIELD = "modal_reserve_cooldown_count"
 MODAL_RESERVE_COOLDOWN_FIELD = "modal_reserve_cooldown_seconds"
+MODAL_RESERVE_COOLDOWN_POOL_FIELD = "modal_reserve_cooldown_pool"
 MODAL_COOLDOWN_PAYLOAD_FIELDS = (
     MODAL_GLOBAL_COOLDOWN_FIELD,
     MODAL_RESERVE_COOLDOWN_COUNT_FIELD,
     MODAL_RESERVE_COOLDOWN_FIELD,
+    MODAL_RESERVE_COOLDOWN_POOL_FIELD,
 )
 """Payload fields carrying the workflow's per-slot Modal cooldown policy.
 
@@ -277,6 +279,21 @@ MODAL_INHERITED_PAYLOAD_FIELDS = (
     *MODAL_COOLDOWN_PAYLOAD_FIELDS,
 )
 
+MODAL_WORKER_POOL_MAPPED = "mapped"
+MODAL_WORKER_POOL_COMFY = "comfy"
+MODAL_WORKER_POOL_LLM = "llm"
+MODAL_WORKER_POOLS = (
+    MODAL_WORKER_POOL_MAPPED,
+    MODAL_WORKER_POOL_COMFY,
+    MODAL_WORKER_POOL_LLM,
+)
+"""Independent Modal worker-slot namespaces.
+
+Mapped lanes use ``worker-pool:slot:N``; ordinary components use
+``worker-pool:<comfy|llm>:slot:N``. Slot N in different pools is a different
+container, so the reserve cooldown must target exactly one pool per prompt.
+"""
+
 MODAL_COOLDOWN_MIN_SECONDS = 2
 MODAL_COOLDOWN_MAX_SECONDS = 3600
 DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS = 60
@@ -321,9 +338,23 @@ class ModalCooldownPolicy:
         if self.reserve_cooldown_count < 0:
             raise ValueError("Modal reserve_cooldown_count must not be negative.")
 
-    def scaledown_window_for_slot(self, slot_index: int | None) -> int:
-        """Return the idle cooldown for one worker slot; unslotted workers are global."""
-        if slot_index is not None and 0 <= slot_index < self.reserve_cooldown_count:
+    def scaledown_window_for_worker(
+        self,
+        *,
+        worker_pool: str | None,
+        slot_index: int | None,
+        reserve_pool: str,
+    ) -> int:
+        """Return one worker's idle cooldown.
+
+        Only the first ``reserve_cooldown_count`` slots of ``reserve_pool`` keep the
+        reserve cooldown; every other worker, including unslotted ones, is global.
+        """
+        if (
+            worker_pool == reserve_pool
+            and slot_index is not None
+            and 0 <= slot_index < self.reserve_cooldown_count
+        ):
             return self.reserve_cooldown_seconds
         return self.global_cooldown_seconds
 

tests/test_configured_remote_planning.py

@@ -666,6 +666,7 @@ def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload(
     assert payload["modal_global_cooldown_seconds"] == 45
     assert payload["modal_reserve_cooldown_count"] == 1
     assert payload["modal_reserve_cooldown_seconds"] == 900
+    assert payload["modal_reserve_cooldown_pool"] == "comfy"
     assert summary.remote_configurations == [
         {
             "configuration_id": "10",

tests/test_execution_scheduling.py

@@ -1202,9 +1202,39 @@ def test_speculative_prewarm_target_keeps_modal_cooldown_policy(
         "modal_global_cooldown_seconds": 45,
         "modal_reserve_cooldown_count": 1,
         "modal_reserve_cooldown_seconds": 900,
+        "modal_reserve_cooldown_pool": "mapped",
     }
     target = prompt_affinity_planning_module._speculative_prewarm_target_payload(
         {"component_id": "component-1", **cooldown_fields}
     )
 
     assert target == {"component_id": "component-1", **cooldown_fields}
+
+
+@pytest.mark.parametrize(
+    ("component_groups", "mapped_ids", "expected_pool"),
+    [
+        ({"1": "comfy", "1__mapped": "comfy"}, {"1__mapped"}, "mapped"),
+        ({"1": "llm", "2__mapped": "comfy"}, {"2__mapped"}, "mapped"),
+        ({"1": "comfy", "2": "llm"}, set(), "comfy"),
+        ({"1": None}, set(), "comfy"),
+        ({"1": "llm"}, set(), "llm"),
+    ],
+)
+def test_reserve_cooldown_pool_targets_one_pool_per_prompt(
+    prompt_affinity_planning_module: Any,
+    component_groups: dict[str, str | None],
+    mapped_ids: set[str],
+    expected_pool: str,
+) -> None:
+    """Exactly one worker pool per Modal target keeps the reserve cooldown."""
+    payloads = {
+        component_id: {"remote_worker_affinity_group": group}
+        for component_id, group in component_groups.items()
+    }
+
+    pool = prompt_affinity_planning_module._reserve_cooldown_pool(
+        list(component_groups), payloads, mapped_ids
+    )
+
+    assert pool == expected_pool

tests/test_modal_deployment.py

@@ -1250,6 +1250,67 @@ def test_lookup_deployed_remote_engine_applies_reserve_cooldown_by_slot(
     assert result["worker_affinity_key"] == f"worker-pool:slot:{lane_index}"
 
 
+@pytest.mark.parametrize(
+    ("reserve_pool", "affinity_key", "expected_window"),
+    [
+        ("mapped", "worker-pool:slot:0", 900),
+        ("mapped", "worker-pool:comfy:slot:0", 45),
+        ("mapped", "worker-pool:llm:slot:0", 45),
+        ("comfy", "worker-pool:comfy:slot:0", 900),
+        ("comfy", "worker-pool:comfy:slot:2", 45),
+        ("comfy", "worker-pool:slot:0", 45),
+        ("llm", "worker-pool:llm:slot:1", 900),
+        ("llm", "worker-pool:comfy:slot:0", 45),
+    ],
+)
+def test_lookup_deployed_remote_engine_reserves_only_the_prompt_reserve_pool(
+    remote_modal_app_module: Any,
+    modal_deployment_module: Any,
+    monkeypatch: Any,
+    reserve_pool: str,
+    affinity_key: str,
+    expected_window: int,
+) -> None:
+    """Slot N in a different worker pool is a different container and stays global."""
+    observed_options = _install_option_recording_modal(
+        remote_modal_app_module, modal_deployment_module, monkeypatch
+    )
+
+    modal_deployment_module._lookup_deployed_remote_engine(
+        {
+            "component_id": "component-1",
+            **_COOLDOWN_PAYLOAD_FIELDS,
+            "modal_reserve_cooldown_pool": reserve_pool,
+        },
+        affinity_key_override=affinity_key,
+    )
+
+    assert observed_options == [{"scaledown_window": expected_window}]
+
+
+def test_lookup_deployed_remote_engine_static_phase_of_mapped_prompt_is_global(
+    remote_modal_app_module: Any,
+    modal_deployment_module: Any,
+    monkeypatch: Any,
+) -> None:
+    """A hybrid prompt's static comfy worker must not hold a second reserved GPU."""
+    observed_options = _install_option_recording_modal(
+        remote_modal_app_module, modal_deployment_module, monkeypatch
+    )
+
+    result = modal_deployment_module._lookup_deployed_remote_engine(
+        {
+            "component_id": "1",
+            "remote_worker_affinity_group": "comfy",
+            **_COOLDOWN_PAYLOAD_FIELDS,
+            "modal_reserve_cooldown_pool": "mapped",
+        }
+    )
+
+    assert result["worker_affinity_key"] == "worker-pool:comfy:slot:0"
+    assert observed_options == [{"scaledown_window": 45}]
+
+
 def test_lookup_deployed_remote_engine_uses_global_cooldown_for_unslotted_keys(
     remote_modal_app_module: Any,
     modal_deployment_module: Any,