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>
350 lines
13 KiB
Python
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",
|
|
]
|