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,