Files
ttulttul-ComfyUI-Modal/remote_configurations.py
T
Ken SimpsonandClaude Opus 5.5 344cfb0142 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 <noreply@anthropic.com>
2026-11-05 12:09:43 -07:00

350 lines
13 KiB
Python

"""Workflow-declared remote execution configuration models."""
from __future__ import annotations
import logging
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass
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]$")
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class RemoteExecutionConfiguration(ABC):
"""Describe one workflow-declared remote-execution configuration value."""
configuration_id: str
display_name: str
def __post_init__(self) -> None:
"""Reject ambiguous identities before provider discovery begins."""
if not self.configuration_id.strip():
raise ValueError("Remote execution configuration_id must not be empty.")
if not self.display_name.strip():
raise ValueError(
"Remote execution configuration display_name must not be empty."
)
@abstractmethod
def to_safe_dict(self) -> dict[str, Any]:
"""Return credential-free workflow and diagnostic metadata."""
@dataclass(frozen=True)
class RemoteConfiguration(RemoteExecutionConfiguration, ABC):
"""Describe one workflow-declared pool of remote execution capacity."""
@property
@abstractmethod
def provider(self) -> ExecutionProvider:
"""Return the provider represented by this configuration."""
@property
@abstractmethod
def capacity_limit(self) -> int:
"""Return the maximum concurrent worker capacity contributed by this pool."""
@dataclass(frozen=True)
class StorageBackingConfiguration(RemoteExecutionConfiguration, ABC):
"""Describe storage shared by one or more remote execution targets."""
@property
@abstractmethod
def storage_provider(self) -> str:
"""Return the stable storage-provider identifier."""
@dataclass(frozen=True)
class R2StorageBackingConfiguration(StorageBackingConfiguration):
"""Reference controller-held credentials for one Cloudflare R2 bucket."""
account_id: str
bucket: str
credential_id: str
jurisdiction: str = "default"
key_prefix: str = "comfy-modal-cache/v1/blobs/sha256"
write_back_mode: str = "async"
def __post_init__(self) -> None:
"""Validate non-secret R2 workflow metadata."""
super().__post_init__()
if len(self.account_id) != 32 or any(
character not in "0123456789abcdefABCDEF" for character in self.account_id
):
raise ValueError(
"Cloudflare R2 account ID must be 32 hexadecimal characters; "
"use Login on the R2 Storage Configuration node."
)
if not _R2_BUCKET_PATTERN.fullmatch(self.bucket):
raise ValueError(
"Cloudflare R2 bucket must contain 3-63 lowercase letters, digits, "
"or hyphens and begin and end with a letter or digit."
)
if not self.credential_id.strip():
raise ValueError("Cloudflare R2 credential_id must not be empty.")
if self.jurisdiction not in {"default", "eu", "fedramp", "us"}:
raise ValueError("Cloudflare R2 jurisdiction is not supported.")
if self.write_back_mode not in {"async", "off", "sync"}:
raise ValueError("Cloudflare R2 write-back mode is not supported.")
if self.write_back_mode == "sync":
logger.warning(
"Legacy synchronous R2 write-back configuration %s now uses idle background write-back.",
self.configuration_id,
)
object.__setattr__(self, "write_back_mode", "async")
@property
def storage_provider(self) -> str:
"""Return the Cloudflare R2 provider identifier."""
return "cloudflare_r2"
def to_safe_dict(self) -> dict[str, Any]:
"""Return R2 metadata without controller credentials."""
return {
"configuration_id": self.configuration_id,
"display_name": self.display_name,
"configuration_kind": "storage",
"storage_provider": self.storage_provider,
"account_id": self.account_id,
"bucket": self.bucket,
"jurisdiction": self.jurisdiction,
"key_prefix": self.key_prefix,
"write_back_mode": self.write_back_mode,
}
@dataclass(frozen=True)
class ModalRemoteConfiguration(RemoteConfiguration):
"""Declare a Modal GPU type and its maximum concurrent container capacity."""
gpu_type: str
instance_count: int = 1
cooldown: ModalCooldownPolicy = ModalCooldownPolicy()
def __post_init__(self) -> None:
"""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:
"""Return the Modal provider identifier."""
return ExecutionProvider.MODAL
@property
def capacity_limit(self) -> int:
"""Return the configured Modal container ceiling."""
return self.instance_count
def to_safe_dict(self) -> dict[str, Any]:
"""Return credential-free Modal configuration metadata."""
return {
"configuration_id": self.configuration_id,
"display_name": self.display_name,
"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,
}
@dataclass(frozen=True)
class VastRemoteConfiguration(RemoteConfiguration):
"""Declare one Vast.ai marketplace capacity pool."""
profile: VastResourceProfile
def __post_init__(self) -> None:
"""Keep the wrapper identity aligned with the underlying Vast profile."""
super().__post_init__()
if self.profile.profile_id != self.configuration_id:
raise ValueError(
"Vast profile_id must match its remote configuration_id."
)
if self.profile.profile_name != self.display_name:
raise ValueError(
"Vast profile_name must match its remote configuration display_name."
)
@property
def provider(self) -> ExecutionProvider:
"""Return the Vast.ai provider identifier."""
return ExecutionProvider.VAST
@property
def capacity_limit(self) -> int:
"""Return the maximum number of managed Vast.ai instances."""
return self.profile.maximum_instances
def to_safe_dict(self) -> dict[str, Any]:
"""Return credential-free Vast configuration metadata."""
return {
"configuration_id": self.configuration_id,
"display_name": self.display_name,
"provider": self.provider.value,
"maximum_instances": self.profile.maximum_instances,
"search": self.profile.search_payload(),
}
@dataclass(frozen=True)
class SshRemoteConfiguration(RemoteConfiguration):
"""Declare one SSH-accessible Docker host and its worker capacity."""
host: SshHostConfig
def __post_init__(self) -> None:
"""Keep workflow and host identities aligned."""
super().__post_init__()
if self.host.environment_id != self.configuration_id:
raise ValueError(
"SSH environment_id must match its remote configuration_id."
)
if self.host.display_name != self.display_name:
raise ValueError(
"SSH display_name must match its remote configuration display_name."
)
if self.host.maximum_workers <= 0:
raise ValueError("SSH maximum_workers must be positive in a workflow.")
@property
def provider(self) -> ExecutionProvider:
"""Return the SSH Docker provider identifier."""
return ExecutionProvider.SSH_DOCKER
@property
def capacity_limit(self) -> int:
"""Return the configured worker-container ceiling for this host."""
return self.host.maximum_workers
def to_safe_dict(self) -> dict[str, Any]:
"""Return credential-free SSH configuration metadata."""
return {
"configuration_id": self.configuration_id,
"display_name": self.display_name,
"provider": self.provider.value,
"host": self.host.to_dict(),
}
@dataclass(frozen=True)
class RemoteConfigurationSet:
"""Hold the authoritative capacity and storage declaration for a workflow."""
configurations: tuple[RemoteExecutionConfiguration, ...]
def __post_init__(self) -> None:
"""Require at least one uniquely identified and named configuration."""
if not self.configurations:
raise ValueError(
"Remote Execution Configurator requires at least one configuration."
)
configuration_ids = [
configuration.configuration_id for configuration in self.configurations
]
if len(configuration_ids) != len(set(configuration_ids)):
raise ValueError("Remote configuration IDs must be unique.")
normalized_names = [
configuration.display_name.casefold()
for configuration in self.configurations
]
if len(normalized_names) != len(set(normalized_names)):
raise ValueError("Remote configuration names must be unique.")
if not any(
isinstance(configuration, RemoteConfiguration)
for configuration in self.configurations
):
raise ValueError(
"Remote Execution Configurator requires at least one Modal, Vast.ai, "
"or SSH capacity configuration."
)
modal_gpu_types = [
configuration.gpu_type
for configuration in self.configurations
if isinstance(configuration, ModalRemoteConfiguration)
]
if len(modal_gpu_types) != len(set(modal_gpu_types)):
raise ValueError(
"Each Modal GPU type may appear in only one remote configuration."
)
ssh_targets = [
configuration.host.ssh_target
for configuration in self.configurations
if isinstance(configuration, SshRemoteConfiguration)
]
if len(ssh_targets) != len(set(ssh_targets)):
raise ValueError(
"Each SSH destination may appear in only one remote configuration."
)
storage_providers = [
configuration.storage_provider
for configuration in self.configurations
if isinstance(configuration, StorageBackingConfiguration)
]
if len(storage_providers) != len(set(storage_providers)):
raise ValueError(
"Each storage provider may appear only once in a Remote Execution "
"Configurator."
)
@property
def capacity_configurations(self) -> tuple[RemoteConfiguration, ...]:
"""Return only configurations that contribute execution capacity."""
return tuple(
configuration
for configuration in self.configurations
if isinstance(configuration, RemoteConfiguration)
)
@property
def storage_configurations(self) -> tuple[StorageBackingConfiguration, ...]:
"""Return only shared storage backing configurations."""
return tuple(
configuration
for configuration in self.configurations
if isinstance(configuration, StorageBackingConfiguration)
)
def to_safe_list(self) -> list[dict[str, Any]]:
"""Return ordered credential-free configuration metadata."""
return [
configuration.to_safe_dict()
for configuration in self.configurations
]
__all__ = [
"ModalRemoteConfiguration",
"RemoteConfiguration",
"RemoteConfigurationSet",
"RemoteExecutionConfiguration",
"R2StorageBackingConfiguration",
"SshRemoteConfiguration",
"StorageBackingConfiguration",
"VastRemoteConfiguration",
]