From 344cfb0142b5fe0f09830271070682a835695c8e Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Thu, 5 Nov 2026 12:09:43 -0700 Subject: [PATCH 1/3] Add global and reserve cooldown settings to Modal Configuration The Modal Configuration node gains three advanced widgets: - global_cooldown (default 60s): idle scaledown for every Modal worker. - reserve_cooldown_count (default 0): how many worker slots use the reserve cooldown; must not exceed instance_count. - reserve_cooldown (default 300s): idle scaledown for the reserved slots. The policy travels on remote payloads and is applied per call with Cls.with_options(scaledown_window=...), keyed by the worker-pool slot parsed from the affinity key. Slots below reserve_cooldown_count stay warm longer; other slots and unslotted keys use the global cooldown. Local-gap pools still keep at least the local keepalive window. The fields do not shape the runtime fingerprint, but derived payloads (warmups, prewarm targets, mapped phases) inherit them so they land in the same container pool. Co-Authored-By: Claude Opus 5.5 --- README.md | 10 ++ docs/LEARNINGS.md | 16 +++ execution_assignment_runtime.py | 1 + prompt_affinity_planning.py | 6 +- remote/modal_deployment.py | 53 ++++++++-- remote_configuration_nodes.py | 85 ++++++++++++++- remote_configurations.py | 13 ++- settings.py | 94 ++++++++++++++++- tests/test_configured_remote_planning.py | 9 ++ tests/test_execution_scheduling.py | 16 +++ tests/test_mapped_execution.py | 6 ++ tests/test_modal_deployment.py | 129 +++++++++++++++++++++++ tests/test_remote_configuration_nodes.py | 86 +++++++++++++++ 13 files changed, 506 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index ada74be..dfa2664 100644 --- a/README.md +++ b/README.md @@ -138,6 +138,16 @@ Each ComfyUI installation gets its own persistent Modal app name, derived from a Each **Modal Configuration** selects a GPU type and maximum concurrent container count. Every GPU target gets its own deployed app, so several configurations can contribute different Modal GPU types to the same plan. First-run auto-deploy is enabled by default and also replaces stale or incompatible deployments; after upgrading this node pack, redeploy once (or leave `COMFY_MODAL_AUTO_DEPLOY=true`) so remote workers pick up the new code. Warm containers reuse loaded models across compatible requests and scale down to zero after an idle window (`COMFY_MODAL_SCALEDOWN_WINDOW`, default 600 seconds). +The Modal Configuration node's advanced settings control how long idle workers stay warm: + +| 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` | `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. + ### Self-Hosted SSH Hosts Self-hosted execution runs the same workers on machines you control, reached over SSH and managed through Docker. Add one **SSH Configuration** per machine and connect it to the configurator. The extension deliberately stores no passwords or private-key paths — authentication, jump hosts, ports, and key selection stay in your normal SSH agent and `~/.ssh/config`. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index c6448f4..da53606 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -766,3 +766,19 @@ generation streams with an empty `modal.exception.RemoteError`. Fixes: copy `settings.MODAL_DEPLOYMENT_PAYLOAD_FIELDS` through `modal_deployment_payload_fields()`; add any new fingerprint-shaping payload field there. + +## Per-slot Modal cooldowns (reserve vs. global) + +Modal's `scaledown_window` is an autoscaler setting, not a per-container one, so +"keep N workers warm longer" is implemented on the worker-pool slots that already +give each `worker-pool:...slot:N` affinity key its own parametrized container pool. +`_lookup_deployed_remote_engine` parses the slot from the affinity key and calls +`Cls.with_options(scaledown_window=...)`: slots `0..reserve_cooldown_count-1` get +`reserve_cooldown`, every other slot (and unslotted keys such as LLM recovery +workers) gets `global_cooldown`. Because it is per call, the cooldown fields are +deliberately *not* fingerprint-shaping (they are kept out of +`_settings_for_payload`), yet derived payloads must still inherit them so warmups +and mapped phases hit the same pool as the real call. They live in +`settings.MODAL_COOLDOWN_PAYLOAD_FIELDS`, and `modal_deployment_payload_fields()` +now copies `MODAL_INHERITED_PAYLOAD_FIELDS` (deployment + cooldown). + diff --git a/execution_assignment_runtime.py b/execution_assignment_runtime.py index 29c6342..6d8c1f1 100644 --- a/execution_assignment_runtime.py +++ b/execution_assignment_runtime.py @@ -225,6 +225,7 @@ def _configured_provider_metadata( "remote_configuration_name": configuration.display_name, "modal_gpu": configuration.gpu_type, "modal_max_containers": configuration.instance_count, + **configuration.cooldown.to_payload_fields(), } if isinstance(configuration, SshRemoteConfiguration): host = execution_plan.ssh_hosts_by_id[configuration.configuration_id] diff --git a/prompt_affinity_planning.py b/prompt_affinity_planning.py index 18133f1..2026f4e 100644 --- a/prompt_affinity_planning.py +++ b/prompt_affinity_planning.py @@ -24,7 +24,7 @@ if __package__: _node_output_refs, ) from .remote_plan_types import InputTarget, LinkedOutputRef - from .settings import MODAL_DEPLOYMENT_PAYLOAD_FIELDS + from .settings import MODAL_INHERITED_PAYLOAD_FIELDS 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 +41,7 @@ else: # pragma: no cover - flat import inside the Modal container. _node_output_refs, ) from remote_plan_types import InputTarget, LinkedOutputRef - from settings import MODAL_DEPLOYMENT_PAYLOAD_FIELDS + from settings import MODAL_INHERITED_PAYLOAD_FIELDS logger = logging.getLogger(__name__) @@ -50,7 +50,7 @@ _SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset( { "component_id", "prompt_id", - *MODAL_DEPLOYMENT_PAYLOAD_FIELDS, + *MODAL_INHERITED_PAYLOAD_FIELDS, "execution_provider", "execution_environment_id", "remote_worker_affinity_group", diff --git a/remote/modal_deployment.py b/remote/modal_deployment.py index 3a9d0b0..198a3b1 100644 --- a/remote/modal_deployment.py +++ b/remote/modal_deployment.py @@ -12,6 +12,7 @@ import inspect import logging import os from pathlib import Path +import re import shutil import subprocess import sys @@ -26,6 +27,7 @@ from ..runtime_environment import ( build_remote_runtime_identity, ) from ..settings import ( + ModalCooldownPolicy, ModalSyncSettings, get_settings, modal_deployment_app_name, @@ -52,6 +54,7 @@ _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+)$") @dataclass(frozen=True) @@ -196,6 +199,38 @@ 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) + if match is None: + return None + return int(match.group(1)) + +def _scaledown_window_override( + payload: Mapping[str, Any], + settings: ModalSyncSettings, + worker_affinity_key: str, +) -> 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. + """ + 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) + ) + if bool(payload.get("remote_local_gap_pool")): + scaledown_window = max( + scaledown_window + if scaledown_window is not None + else int(settings.scaledown_window_seconds), + int(settings.local_gap_keepalive_seconds), + ) + return scaledown_window + def _modal_lookup_error_types() -> tuple[type[BaseException], ...]: """Return Modal exception types that indicate lookup or hydration failure.""" if modal is None: @@ -359,19 +394,17 @@ def _lookup_deployed_remote_engine( gpu_snapshot_enabled, ) remote_cls = modal.Cls.from_name(deployment_app_name, "RemoteEngine") - if bool(payload.get("remote_local_gap_pool")) and hasattr(remote_cls, "with_options"): - local_gap_scaledown_seconds = max( - int(settings.scaledown_window_seconds), - int(settings.local_gap_keepalive_seconds), - ) - remote_cls = remote_cls.with_options( - scaledown_window=local_gap_scaledown_seconds, - ) + scaledown_window = _scaledown_window_override( + payload, settings, worker_affinity_key + ) + if scaledown_window is not None and hasattr(remote_cls, "with_options"): + remote_cls = remote_cls.with_options(scaledown_window=scaledown_window) logger.info( - "Using local-gap Modal pool component=%s worker_affinity=%s scaledown_window=%ds.", + "Using Modal pool component=%s worker_affinity=%s local_gap=%s scaledown_window=%ds.", payload.get("component_id"), worker_affinity_key, - local_gap_scaledown_seconds, + bool(payload.get("remote_local_gap_pool")), + scaledown_window, ) remote_engine_kwargs: dict[str, Any] = { "gpu_snapshot_enabled": gpu_snapshot_enabled, diff --git a/remote_configuration_nodes.py b/remote_configuration_nodes.py index 33be113..8c31908 100644 --- a/remote_configuration_nodes.py +++ b/remote_configuration_nodes.py @@ -19,7 +19,16 @@ if __package__: VastRemoteConfiguration, ) from .remote_hosts import SshHostConfig - from .settings import MODAL_GPU_TYPES, normalize_modal_gpu_selection + from .settings import ( + DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS, + DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT, + DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS, + MODAL_COOLDOWN_MAX_SECONDS, + MODAL_COOLDOWN_MIN_SECONDS, + MODAL_GPU_TYPES, + ModalCooldownPolicy, + normalize_modal_gpu_selection, + ) from .vast_config_node import VastAILeaseConfiguration, profile_from_inputs else: # pragma: no cover - direct ComfyUI loading fallback. from remote_configurations import ( @@ -32,7 +41,16 @@ else: # pragma: no cover - direct ComfyUI loading fallback. VastRemoteConfiguration, ) from remote_hosts import SshHostConfig - from settings import MODAL_GPU_TYPES, normalize_modal_gpu_selection + from settings import ( + DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS, + DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT, + DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS, + MODAL_COOLDOWN_MAX_SECONDS, + MODAL_COOLDOWN_MIN_SECONDS, + MODAL_GPU_TYPES, + ModalCooldownPolicy, + normalize_modal_gpu_selection, + ) from vast_config_node import VastAILeaseConfiguration, profile_from_inputs logger = logging.getLogger(__name__) @@ -185,6 +203,7 @@ class ModalConfiguration(io.ComfyNode): step=1, tooltip="Maximum concurrent Modal containers for this pool.", ), + *_modal_cooldown_inputs(), ], outputs=[RemoteConfigurationType.Output()], hidden=[io.Hidden.unique_id], @@ -200,6 +219,49 @@ class ModalConfiguration(io.ComfyNode): return io.NodeOutput(modal_configuration_from_inputs(configuration_id, inputs)) +def _modal_cooldown_inputs() -> list[io.Input]: + """Return advanced widgets controlling how long idle Modal workers stay warm.""" + return [ + io.Int.Input( + "global_cooldown", + default=DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS, + min=MODAL_COOLDOWN_MIN_SECONDS, + max=MODAL_COOLDOWN_MAX_SECONDS, + step=1, + advanced=True, + tooltip=( + "Seconds an idle Modal worker stays warm before scaling down. " + "Applies to every worker not covered by reserve_cooldown_count." + ), + ), + io.Int.Input( + "reserve_cooldown_count", + default=DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT, + min=0, + max=32, + step=1, + advanced=True, + tooltip=( + "Number of Modal workers (the first N worker slots) that use " + "reserve_cooldown instead of global_cooldown. Must not exceed " + "instance_count." + ), + ), + io.Int.Input( + "reserve_cooldown", + default=DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS, + min=MODAL_COOLDOWN_MIN_SECONDS, + max=MODAL_COOLDOWN_MAX_SECONDS, + step=1, + advanced=True, + tooltip=( + "Seconds the reserved Modal workers stay warm while idle. Set it " + "longer than global_cooldown to keep a small always-warm set." + ), + ), + ] + + class VastConfiguration(io.ComfyNode): """Declare one workflow-scoped Vast.ai marketplace capacity pool.""" @@ -367,6 +429,24 @@ def modal_configuration_from_inputs( display_name=str(inputs.get("configuration_name") or "modal-default").strip(), gpu_type=normalize_modal_gpu_selection(inputs.get("gpu_type", "RTX-PRO-6000")), instance_count=int(inputs.get("instance_count", 1)), + cooldown=modal_cooldown_policy_from_inputs(inputs), + ) + + +def modal_cooldown_policy_from_inputs( + inputs: Mapping[str, Any], +) -> ModalCooldownPolicy: + """Build the Modal cooldown policy, defaulting widgets absent from older workflows.""" + return ModalCooldownPolicy( + global_cooldown_seconds=int( + inputs.get("global_cooldown", DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS) + ), + reserve_cooldown_count=int( + inputs.get("reserve_cooldown_count", DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT) + ), + reserve_cooldown_seconds=int( + inputs.get("reserve_cooldown", DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS) + ), ) @@ -653,6 +733,7 @@ __all__ = [ "compile_remote_configuration_set", "configuration_set_from_values", "modal_configuration_from_inputs", + "modal_cooldown_policy_from_inputs", "r2_storage_configuration_from_inputs", "ssh_configuration_from_inputs", "vast_configuration_from_inputs", diff --git a/remote_configurations.py b/remote_configurations.py index b042b93..a49ec52 100644 --- a/remote_configurations.py +++ b/remote_configurations.py @@ -11,10 +11,12 @@ from typing import Any if __package__: from .execution_environments import ExecutionProvider from .remote_hosts import SshHostConfig + from .settings import ModalCooldownPolicy from .vast_models import VastResourceProfile else: # pragma: no cover - direct ComfyUI loading fallback. from execution_environments import ExecutionProvider from remote_hosts import SshHostConfig + from settings import ModalCooldownPolicy from vast_models import VastResourceProfile _R2_BUCKET_PATTERN = re.compile(r"^[a-z0-9][a-z0-9-]{1,61}[a-z0-9]$") @@ -131,14 +133,20 @@ class ModalRemoteConfiguration(RemoteConfiguration): gpu_type: str instance_count: int = 1 + cooldown: ModalCooldownPolicy = ModalCooldownPolicy() def __post_init__(self) -> None: - """Validate Modal capacity fields.""" + """Validate Modal capacity and cooldown fields.""" super().__post_init__() if not self.gpu_type.strip(): raise ValueError("Modal gpu_type must not be empty.") if self.instance_count <= 0: raise ValueError("Modal instance_count must be positive.") + if self.cooldown.reserve_cooldown_count > self.instance_count: + raise ValueError( + f"Modal reserve_cooldown_count ({self.cooldown.reserve_cooldown_count}) " + f"must not exceed instance_count ({self.instance_count})." + ) @property def provider(self) -> ExecutionProvider: @@ -158,6 +166,9 @@ class ModalRemoteConfiguration(RemoteConfiguration): "provider": self.provider.value, "gpu_type": self.gpu_type, "instance_count": self.instance_count, + "global_cooldown_seconds": self.cooldown.global_cooldown_seconds, + "reserve_cooldown_count": self.cooldown.reserve_cooldown_count, + "reserve_cooldown_seconds": self.cooldown.reserve_cooldown_seconds, } diff --git a/settings.py b/settings.py index 2bb3b37..13041a8 100644 --- a/settings.py +++ b/settings.py @@ -258,16 +258,106 @@ Any payload derived from another (warmups, prewarm targets, mapped phases) must 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_COOLDOWN_PAYLOAD_FIELDS = ( + MODAL_GLOBAL_COOLDOWN_FIELD, + MODAL_RESERVE_COOLDOWN_COUNT_FIELD, + MODAL_RESERVE_COOLDOWN_FIELD, +) +"""Payload fields carrying the workflow's per-slot Modal cooldown policy. + +They are applied per call through ``with_options`` and never change the +deployed app's fingerprint, but derived payloads must still inherit them. +""" + +MODAL_INHERITED_PAYLOAD_FIELDS = ( + *MODAL_DEPLOYMENT_PAYLOAD_FIELDS, + *MODAL_COOLDOWN_PAYLOAD_FIELDS, +) + +MODAL_COOLDOWN_MIN_SECONDS = 2 +MODAL_COOLDOWN_MAX_SECONDS = 3600 +DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS = 60 +DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT = 0 +DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS = 300 + def modal_deployment_payload_fields(payload: Mapping[str, Any]) -> dict[str, Any]: - """Return the deployment-identity fields present on one payload.""" + """Return the deployment-identity and cooldown fields present on one payload.""" return { field_name: payload[field_name] - for field_name in MODAL_DEPLOYMENT_PAYLOAD_FIELDS + for field_name in MODAL_INHERITED_PAYLOAD_FIELDS if payload.get(field_name) is not None } +def _validate_cooldown_seconds(name: str, value: int) -> None: + """Reject a cooldown that Modal's autoscaler would not accept.""" + if not MODAL_COOLDOWN_MIN_SECONDS <= value <= MODAL_COOLDOWN_MAX_SECONDS: + raise ValueError( + f"Modal {name} must be between {MODAL_COOLDOWN_MIN_SECONDS} and " + f"{MODAL_COOLDOWN_MAX_SECONDS} seconds; got {value}." + ) + + +@dataclass(frozen=True) +class ModalCooldownPolicy: + """Assign Modal idle cooldowns to reusable worker-pool slots. + + Slots ``0 .. reserve_cooldown_count - 1`` keep the longer reserve cooldown so + a small set of workers stays warm; every other slot uses the global cooldown. + """ + + global_cooldown_seconds: int = DEFAULT_MODAL_GLOBAL_COOLDOWN_SECONDS + reserve_cooldown_count: int = DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT + reserve_cooldown_seconds: int = DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS + + def __post_init__(self) -> None: + """Validate cooldown bounds and the reserved slot count.""" + _validate_cooldown_seconds("global_cooldown", self.global_cooldown_seconds) + _validate_cooldown_seconds("reserve_cooldown", self.reserve_cooldown_seconds) + 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: + return self.reserve_cooldown_seconds + return self.global_cooldown_seconds + + def to_payload_fields(self) -> dict[str, int]: + """Return the payload fields that carry this policy to remote dispatch.""" + return { + MODAL_GLOBAL_COOLDOWN_FIELD: self.global_cooldown_seconds, + MODAL_RESERVE_COOLDOWN_COUNT_FIELD: self.reserve_cooldown_count, + MODAL_RESERVE_COOLDOWN_FIELD: self.reserve_cooldown_seconds, + } + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> ModalCooldownPolicy | None: + """Return the policy carried by a payload, or None when it carries none.""" + raw_global = payload.get(MODAL_GLOBAL_COOLDOWN_FIELD) + if raw_global is None: + return None + return cls( + global_cooldown_seconds=int(raw_global), + reserve_cooldown_count=int( + payload.get( + MODAL_RESERVE_COOLDOWN_COUNT_FIELD, + DEFAULT_MODAL_RESERVE_COOLDOWN_COUNT, + ) + ), + reserve_cooldown_seconds=int( + payload.get( + MODAL_RESERVE_COOLDOWN_FIELD, + DEFAULT_MODAL_RESERVE_COOLDOWN_SECONDS, + ) + ), + ) + + def _read_path_env(name: str) -> Path | None: """Resolve an environment variable into a path when present.""" value = os.getenv(name) diff --git a/tests/test_configured_remote_planning.py b/tests/test_configured_remote_planning.py index 3958068..992267e 100644 --- a/tests/test_configured_remote_planning.py +++ b/tests/test_configured_remote_planning.py @@ -630,6 +630,9 @@ def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload( "configuration_name": "modal-t4", "gpu_type": "T4", "instance_count": 2, + "global_cooldown": 45, + "reserve_cooldown_count": 1, + "reserve_cooldown": 900, }, }, "99": { @@ -660,6 +663,9 @@ def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload( assert payload["remote_configuration_name"] == "modal-t4" assert payload["modal_gpu"] == "T4" assert payload["modal_max_containers"] == 2 + assert payload["modal_global_cooldown_seconds"] == 45 + assert payload["modal_reserve_cooldown_count"] == 1 + assert payload["modal_reserve_cooldown_seconds"] == 900 assert summary.remote_configurations == [ { "configuration_id": "10", @@ -667,5 +673,8 @@ def test_rewrite_stamps_selected_modal_configuration_on_proxy_payload( "provider": "modal", "gpu_type": "T4", "instance_count": 2, + "global_cooldown_seconds": 45, + "reserve_cooldown_count": 1, + "reserve_cooldown_seconds": 900, } ] diff --git a/tests/test_execution_scheduling.py b/tests/test_execution_scheduling.py index 7ff1bb3..d5731fd 100644 --- a/tests/test_execution_scheduling.py +++ b/tests/test_execution_scheduling.py @@ -1192,3 +1192,19 @@ def test_speculative_prewarm_target_keeps_modal_max_containers( "modal_gpu": "RTX-PRO-6000", "modal_max_containers": 4, } + + +def test_speculative_prewarm_target_keeps_modal_cooldown_policy( + prompt_affinity_planning_module: Any, +) -> None: + """Speculative prewarms must land in the same cooldown pool as the real call.""" + cooldown_fields = { + "modal_global_cooldown_seconds": 45, + "modal_reserve_cooldown_count": 1, + "modal_reserve_cooldown_seconds": 900, + } + target = prompt_affinity_planning_module._speculative_prewarm_target_payload( + {"component_id": "component-1", **cooldown_fields} + ) + + assert target == {"component_id": "component-1", **cooldown_fields} diff --git a/tests/test_mapped_execution.py b/tests/test_mapped_execution.py index 6c6a05b..c41d8f7 100644 --- a/tests/test_mapped_execution.py +++ b/tests/test_mapped_execution.py @@ -2552,6 +2552,9 @@ def test_mapped_phase_payloads_preserve_parent_runtime_fingerprint( "component_id": "1", "modal_gpu": "B300", "modal_max_containers": 4, + "modal_global_cooldown_seconds": 45, + "modal_reserve_cooldown_count": 1, + "modal_reserve_cooldown_seconds": 900, "static_phase": phase_definition, "mapped_phase": phase_definition, "remote_session": {"session_id": "session-1"}, @@ -2585,6 +2588,9 @@ def test_mapped_phase_payloads_preserve_parent_runtime_fingerprint( for derived_payload in derived_payloads: assert derived_payload["modal_gpu"] == "B300" assert derived_payload["modal_max_containers"] == 4 + assert derived_payload["modal_global_cooldown_seconds"] == 45 + assert derived_payload["modal_reserve_cooldown_count"] == 1 + assert derived_payload["modal_reserve_cooldown_seconds"] == 900 assert ( modal_deployment_module._expected_remote_runtime_fingerprint(derived_payload) == parent_fingerprint diff --git a/tests/test_modal_deployment.py b/tests/test_modal_deployment.py index d72d1bb..4633ad1 100644 --- a/tests/test_modal_deployment.py +++ b/tests/test_modal_deployment.py @@ -1178,6 +1178,135 @@ def test_lookup_deployed_remote_engine_extends_local_gap_pool_scaledown( assert observed_kwargs == [result] assert result["worker_affinity_key"] == "worker-pool:slot:0" +def _install_option_recording_modal( + remote_modal_app_module: Any, + modal_deployment_module: Any, + monkeypatch: Any, +) -> list[dict[str, Any]]: + """Install a fake Modal SDK that records ``with_options`` overrides.""" + observed_options: list[dict[str, Any]] = [] + + class FakeRemoteCls: + """Minimal deployed class supporting Modal runtime option overrides.""" + + def with_options(self, **kwargs: Any) -> "FakeRemoteCls": + """Record the per-call autoscaler configuration.""" + observed_options.append(kwargs) + return self + + def __call__(self, **kwargs: Any) -> dict[str, Any]: + """Return the parametrized class instance identity.""" + return kwargs + + class FakeModal: + """Minimal Modal SDK double for deployed class lookup.""" + + class Cls: + """Namespace for deployed class lookups.""" + + @staticmethod + def from_name(app_name: str, class_name: str) -> FakeRemoteCls: + """Return a fresh fake deployed class.""" + del app_name, class_name + return FakeRemoteCls() + + monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal) + monkeypatch.setattr(modal_deployment_module, "modal", FakeModal) + return observed_options + + +_COOLDOWN_PAYLOAD_FIELDS = { + "modal_global_cooldown_seconds": 45, + "modal_reserve_cooldown_count": 2, + "modal_reserve_cooldown_seconds": 900, +} + + +@pytest.mark.parametrize( + ("lane_index", "expected_window"), + [(0, 900), (1, 900), (2, 45), (3, 45)], +) +def test_lookup_deployed_remote_engine_applies_reserve_cooldown_by_slot( + remote_modal_app_module: Any, + modal_deployment_module: Any, + monkeypatch: Any, + lane_index: int, + expected_window: int, +) -> None: + """The first reserve_cooldown_count slots keep the reserve cooldown.""" + observed_options = _install_option_recording_modal( + remote_modal_app_module, modal_deployment_module, monkeypatch + ) + payload = {"component_id": "component-1", **_COOLDOWN_PAYLOAD_FIELDS} + + result = modal_deployment_module._lookup_deployed_remote_engine( + payload, + affinity_key_override=modal_deployment_module._mapped_lane_affinity_key( + payload, lane_index + ), + ) + + assert observed_options == [{"scaledown_window": expected_window}] + assert result["worker_affinity_key"] == f"worker-pool:slot:{lane_index}" + + +def test_lookup_deployed_remote_engine_uses_global_cooldown_for_unslotted_keys( + remote_modal_app_module: Any, + modal_deployment_module: Any, + monkeypatch: Any, +) -> None: + """Ad-hoc affinity keys (e.g. fresh LLM recovery workers) use the global cooldown.""" + 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}, + affinity_key_override="llm-recovery:abc123", + ) + + assert observed_options == [{"scaledown_window": 45}] + + +def test_lookup_deployed_remote_engine_keeps_deployed_cooldown_without_policy( + remote_modal_app_module: Any, + modal_deployment_module: Any, + monkeypatch: Any, +) -> None: + """Payloads without a workflow cooldown policy keep the deployed default.""" + 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"} + ) + + assert observed_options == [] + + +def test_lookup_deployed_remote_engine_local_gap_pool_never_shortens_keepalive( + remote_modal_app_module: Any, + modal_deployment_module: Any, + monkeypatch: Any, +) -> None: + """Local-gap pools keep at least the local keepalive even with a short cooldown.""" + 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", + "remote_local_gap_pool": True, + **_COOLDOWN_PAYLOAD_FIELDS, + "modal_reserve_cooldown_count": 0, + } + ) + + assert observed_options == [{"scaledown_window": 900}] + + def test_lookup_deployed_remote_engine_passes_snapshot_profile_parameter_for_gpu_snapshots( remote_modal_app_module: Any, modal_warmup_module: Any, diff --git a/tests/test_remote_configuration_nodes.py b/tests/test_remote_configuration_nodes.py index 3d41e47..6038567 100644 --- a/tests/test_remote_configuration_nodes.py +++ b/tests/test_remote_configuration_nodes.py @@ -417,3 +417,89 @@ def test_extension_registers_all_remote_configuration_nodes( extension_package.RemoteExecutionConfigurator is module.RemoteExecutionConfigurator ) + + +def test_modal_configuration_exposes_advanced_cooldown_widgets( + remote_configuration_nodes_module: Any, +) -> None: + """The Modal node should expose three advanced cooldown widgets with defaults.""" + schema = remote_configuration_nodes_module.ModalConfiguration.define_schema() + inputs_by_id = {node_input.id: node_input for node_input in schema.inputs} + + expected_defaults = { + "global_cooldown": 60, + "reserve_cooldown_count": 0, + "reserve_cooldown": 300, + } + for input_id, default in expected_defaults.items(): + assert inputs_by_id[input_id].default == default + assert inputs_by_id[input_id].advanced is True + + +def test_modal_configuration_defaults_cooldowns_for_older_workflows( + remote_configuration_nodes_module: Any, +) -> None: + """Workflows saved before the cooldown widgets existed keep working.""" + configuration = remote_configuration_nodes_module.modal_configuration_from_inputs( + "10", + {"configuration_name": "modal", "gpu_type": "T4", "instance_count": 2}, + ) + + assert configuration.cooldown.global_cooldown_seconds == 60 + assert configuration.cooldown.reserve_cooldown_count == 0 + assert configuration.cooldown.reserve_cooldown_seconds == 300 + + +def test_modal_configuration_parses_cooldown_widgets( + remote_configuration_nodes_module: Any, +) -> None: + """Configured cooldown widgets should become the configuration's policy.""" + configuration = remote_configuration_nodes_module.modal_configuration_from_inputs( + "10", + { + "configuration_name": "modal", + "gpu_type": "T4", + "instance_count": 4, + "global_cooldown": 30, + "reserve_cooldown_count": 2, + "reserve_cooldown": 1200, + }, + ) + + assert configuration.cooldown.global_cooldown_seconds == 30 + assert configuration.cooldown.reserve_cooldown_count == 2 + assert configuration.cooldown.reserve_cooldown_seconds == 1200 + assert configuration.to_safe_dict()["reserve_cooldown_count"] == 2 + + +def test_modal_configuration_rejects_reserve_count_above_instance_count( + remote_configuration_nodes_module: Any, +) -> None: + """More reserved workers than the pool can hold is a configuration error.""" + with pytest.raises(ValueError, match="reserve_cooldown_count"): + remote_configuration_nodes_module.modal_configuration_from_inputs( + "10", + { + "configuration_name": "modal", + "gpu_type": "T4", + "instance_count": 1, + "reserve_cooldown_count": 2, + }, + ) + + +@pytest.mark.parametrize("cooldown", [0, 1, 3601]) +def test_modal_configuration_rejects_out_of_range_cooldowns( + remote_configuration_nodes_module: Any, + cooldown: int, +) -> None: + """Cooldowns outside Modal's accepted scaledown range should fail early.""" + with pytest.raises(ValueError, match="global_cooldown"): + remote_configuration_nodes_module.modal_configuration_from_inputs( + "10", + { + "configuration_name": "modal", + "gpu_type": "T4", + "global_cooldown": cooldown, + }, + ) From 62f0a13acc69628b283809ce2cc1f78d067335a4 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Thu, 5 Nov 2026 12:27:26 -0700 Subject: [PATCH 2/3] Fix local fallback executor on current ComfyUI cache args Current ComfyUI's PromptExecutor.execute_async indexes cache_args["ram"] and cache_args["ram_inactive"] on every run. Local fallback execution built its nested executor without cache_args, so every locally executed subgraph failed with "'NoneType' object is not subscriptable". Pass zeroed classic-cache args, and restore ComfyUI's module-global RAM cache release state after the nested run so the host executor keeps its RAM-pressure eviction. Co-Authored-By: Claude Opus 5.5 --- docs/LEARNINGS.md | 11 +++ remote/local_execution.py | 65 ++++++++++++++++- tests/test_local_nested_executor.py | 105 ++++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 2 deletions(-) create mode 100644 tests/test_local_nested_executor.py diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index da53606..16ba85f 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -782,3 +782,14 @@ and mapped phases hit the same pool as the real call. They live in `settings.MODAL_COOLDOWN_PAYLOAD_FIELDS`, and `modal_deployment_payload_fields()` now copies `MODAL_INHERITED_PAYLOAD_FIELDS` (deployment + cooldown). + +## Nested local PromptExecutor needs explicit RAM cache args + +ComfyUI `5aa5ccc9` made `PromptExecutor.execute_async` read +`cache_args["ram"]` and `cache_args["ram_inactive"]` unconditionally, then call +`comfy.memory_management.set_ram_cache_release_state(...)`, which sets +*module globals*. Local fallback execution (`execution_mode=local`) builds a +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. diff --git a/remote/local_execution.py b/remote/local_execution.py index 94c85fa..3c0f4af 100644 --- a/remote/local_execution.py +++ b/remote/local_execution.py @@ -1047,6 +1047,67 @@ def _execute_subgraph_with_mapping( return tuple(outputs) +_NESTED_PROMPT_EXECUTOR_CACHE_ARGS: dict[str, float] = { + "lru": 0, + "ram": 0.0, + "ram_inactive": 0.0, +} +_MEMORY_MANAGEMENT_MODULE_NAME = "comfy.memory_management" + + +@contextmanager +def _preserved_ram_cache_release_state() -> Iterator[None]: + """Restore ComfyUI's global RAM-cache release state after a nested executor run. + + ``PromptExecutor.execute`` installs its own cache release callback and headroom + as module globals; a nested executor must not leave the host executor's + RAM-pressure eviction disabled once it returns. + """ + memory_management = sys.modules.get(_MEMORY_MANAGEMENT_MODULE_NAME) + setter = getattr(memory_management, "set_ram_cache_release_state", None) + if setter is None: + yield + return + saved_callback = getattr(memory_management, "extra_ram_release_callback", None) + saved_headroom = getattr(memory_management, "RAM_CACHE_HEADROOM", 0) + try: + yield + finally: + setter(saved_callback, saved_headroom) + logger.debug( + "Restored host RAM cache release state headroom=%s callback_set=%s.", + saved_headroom, + saved_callback is not None, + ) + + +def _run_nested_prompt_executor( + execution: Any, + *, + prompt: dict[str, Any], + prompt_id: str, + extra_data: dict[str, Any], + execute_outputs: list[str], +) -> Any: + """Run a headless classic-cache PromptExecutor inside the host ComfyUI process. + + Current ComfyUI reads ``cache_args["ram"]`` and ``["ram_inactive"]`` on every + execute, so the nested executor needs explicit values even with classic caching. + """ + executor = execution.PromptExecutor( + _NullPromptServer(), + cache_args=dict(_NESTED_PROMPT_EXECUTOR_CACHE_ARGS), + ) + with _preserved_ram_cache_release_state(): + executor.execute( + prompt=prompt, + prompt_id=prompt_id, + extra_data=extra_data, + execute_outputs=execute_outputs, + ) + return executor + + def _execute_subgraph_prompt( payload: dict[str, Any], hydrated_inputs: dict[str, Any], @@ -1118,14 +1179,14 @@ def _execute_subgraph_prompt( _validate_required_prompt_inputs(prompt, resolved_node_mapping) with _temporary_node_mapping(node_mapping): - executor = execution.PromptExecutor(_NullPromptServer()) execution_started_at = time.perf_counter() logger.info( "Starting PromptExecutor for remote subgraph %s with execute targets %s.", payload.get("component_id"), normalized_payload.get("execute_node_ids", []), ) - executor.execute( + executor = _run_nested_prompt_executor( + execution, prompt=prompt, prompt_id=str( payload.get("prompt_id") diff --git a/tests/test_local_nested_executor.py b/tests/test_local_nested_executor.py new file mode 100644 index 0000000..5d4d092 --- /dev/null +++ b/tests/test_local_nested_executor.py @@ -0,0 +1,105 @@ +"""Tests for the nested PromptExecutor used by local fallback execution.""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace +from typing import Any + + +def _install_fake_memory_management(monkeypatch: Any) -> ModuleType: + """Install a minimal stand-in for ComfyUI's RAM cache release globals.""" + memory_management = ModuleType("comfy.memory_management") + memory_management.extra_ram_release_callback = None + memory_management.RAM_CACHE_HEADROOM = 0 + + def set_ram_cache_release_state(callback: Any, headroom: int) -> None: + """Mirror ComfyUI's module-global setter.""" + memory_management.extra_ram_release_callback = callback + memory_management.RAM_CACHE_HEADROOM = max(0, int(headroom)) + + memory_management.set_ram_cache_release_state = set_ram_cache_release_state + monkeypatch.setitem(sys.modules, "comfy.memory_management", memory_management) + return memory_management + + +def _fake_execution_module(memory_management: ModuleType) -> SimpleNamespace: + """Return an execution module whose executor behaves like current ComfyUI.""" + + class FakePromptExecutor: + """Require RAM cache args and clobber the global release state like ComfyUI.""" + + def __init__( + self, + server: Any, + cache_type: Any = False, + cache_args: Any = None, + ) -> None: + """Record construction arguments.""" + self.server = server + self.cache_type = cache_type + self.cache_args = cache_args + self.executed_outputs: list[str] = [] + + def execute( + self, + *, + prompt: dict[str, Any], + prompt_id: str, + extra_data: dict[str, Any], + execute_outputs: list[str], + ) -> None: + """Fail like ComfyUI on missing cache args, then reset global state.""" + del prompt, prompt_id, extra_data + headroom = int(self.cache_args["ram"] * (1024**3)) + float(self.cache_args["ram_inactive"]) + memory_management.set_ram_cache_release_state(None, headroom) + self.executed_outputs = execute_outputs + + return SimpleNamespace(PromptExecutor=FakePromptExecutor) + + +def test_nested_executor_supplies_ram_cache_args( + local_execution_module: Any, + monkeypatch: Any, +) -> None: + """Current ComfyUI indexes cache_args['ram'] on every execute.""" + memory_management = _install_fake_memory_management(monkeypatch) + + executor = local_execution_module._run_nested_prompt_executor( + _fake_execution_module(memory_management), + prompt={}, + prompt_id="prompt-1", + extra_data={}, + execute_outputs=["9"], + ) + + assert executor.cache_type is False + assert executor.cache_args == {"lru": 0, "ram": 0.0, "ram_inactive": 0.0} + assert executor.executed_outputs == ["9"] + + +def test_nested_executor_restores_host_ram_release_state( + local_execution_module: Any, + monkeypatch: Any, +) -> None: + """The host executor's RAM-pressure eviction must survive a nested run.""" + memory_management = _install_fake_memory_management(monkeypatch) + + def host_release(target: int, free_active: bool = False) -> int: + """Stand in for the host cache's RAM release callback.""" + del target, free_active + return 0 + + memory_management.set_ram_cache_release_state(host_release, 8 * 1024**3) + + local_execution_module._run_nested_prompt_executor( + _fake_execution_module(memory_management), + prompt={}, + prompt_id="prompt-1", + extra_data={}, + execute_outputs=["9"], + ) + + assert memory_management.extra_ram_release_callback is host_release + assert memory_management.RAM_CACHE_HEADROOM == 8 * 1024**3 From 25426e2ef68810a4c7482dd426320cff07080d8a Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Thu, 5 Nov 2026 13:44:06 -0700 Subject: [PATCH 3/3] 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 --- README.md | 2 +- docs/LEARNINGS.md | 13 +++ prompt_affinity_planning.py | 105 ++++++++++++++++++++++- prompt_interception.py | 7 ++ remote/modal_deployment.py | 34 +++++--- settings.py | 37 +++++++- tests/test_configured_remote_planning.py | 1 + tests/test_execution_scheduling.py | 30 +++++++ tests/test_modal_deployment.py | 61 +++++++++++++ 9 files changed, 274 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index dfa2664..59ed667 100644 --- a/README.md +++ b/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. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 16ba85f..2ea185e 100644 --- a/docs/LEARNINGS.md +++ b/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. + diff --git a/prompt_affinity_planning.py b/prompt_affinity_planning.py index 2026f4e..eefe8a7 100644 --- a/prompt_affinity_planning.py +++ b/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 "", + environment_component_ids, + ) + + def _configure_speculative_affinity_prewarm_payloads( *, rewritten_prompt: dict[str, Any], diff --git a/prompt_interception.py b/prompt_interception.py index e1a5fbd..ff827fd 100644 --- a/prompt_interception.py +++ b/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, diff --git a/remote/modal_deployment.py b/remote/modal_deployment.py index 198a3b1..937b143 100644 --- a/remote/modal_deployment.py +++ b/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[a-z]+):)?slot:(?P\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( diff --git a/settings.py b/settings.py index 13041a8..3ecdac7 100644 --- a/settings.py +++ b/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::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 diff --git a/tests/test_configured_remote_planning.py b/tests/test_configured_remote_planning.py index 992267e..2c4e8cd 100644 --- a/tests/test_configured_remote_planning.py +++ b/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", diff --git a/tests/test_execution_scheduling.py b/tests/test_execution_scheduling.py index d5731fd..9297b11 100644 --- a/tests/test_execution_scheduling.py +++ b/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 diff --git a/tests/test_modal_deployment.py b/tests/test_modal_deployment.py index 4633ad1..120b938 100644 --- a/tests/test_modal_deployment.py +++ b/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,