diff --git a/README.md b/README.md index ada74be..59ed667 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 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. + ### 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..2ea185e 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -766,3 +766,43 @@ 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). + + +## 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. + +### 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/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..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_DEPLOYMENT_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_DEPLOYMENT_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__) @@ -50,7 +64,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", @@ -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/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/remote/modal_deployment.py b/remote/modal_deployment.py index 3a9d0b0..937b143 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,9 @@ from ..runtime_environment import ( build_remote_runtime_identity, ) from ..settings import ( + MODAL_RESERVE_COOLDOWN_POOL_FIELD, + MODAL_WORKER_POOL_MAPPED, + ModalCooldownPolicy, ModalSyncSettings, get_settings, modal_deployment_app_name, @@ -52,6 +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"^worker-pool:(?:(?P[a-z]+):)?slot:(?P\d+)$" +) @dataclass(frozen=True) @@ -196,6 +203,48 @@ 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_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, None + return match.group("pool") or MODAL_WORKER_POOL_MAPPED, int(match.group("slot")) + +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 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: + 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( + 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 +408,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..3ecdac7 100644 --- a/settings.py +++ b/settings.py @@ -258,16 +258,137 @@ 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_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. + +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_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 +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_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 + + 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..2c4e8cd 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,10 @@ 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 payload["modal_reserve_cooldown_pool"] == "comfy" assert summary.remote_configurations == [ { "configuration_id": "10", @@ -667,5 +674,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..9297b11 100644 --- a/tests/test_execution_scheduling.py +++ b/tests/test_execution_scheduling.py @@ -1192,3 +1192,49 @@ 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, + "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_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 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..120b938 100644 --- a/tests/test_modal_deployment.py +++ b/tests/test_modal_deployment.py @@ -1178,6 +1178,196 @@ 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}" + + +@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, + 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, + }, + )