Files
ttulttul-ComfyUI-Modal/remote_configuration_nodes.py
T

660 lines
24 KiB
Python

"""ComfyUI v3 nodes for workflow-scoped remote execution configuration."""
from __future__ import annotations
import logging
from collections.abc import Mapping
from typing import Any
from comfy_api.latest import _io as io
if __package__:
from .remote_configurations import (
ModalRemoteConfiguration,
RemoteConfiguration,
RemoteConfigurationSet,
RemoteExecutionConfiguration,
R2StorageBackingConfiguration,
SshRemoteConfiguration,
VastRemoteConfiguration,
)
from .remote_hosts import SshHostConfig
from .settings import MODAL_GPU_TYPES, 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 (
ModalRemoteConfiguration,
RemoteConfiguration,
RemoteConfigurationSet,
RemoteExecutionConfiguration,
R2StorageBackingConfiguration,
SshRemoteConfiguration,
VastRemoteConfiguration,
)
from remote_hosts import SshHostConfig
from settings import MODAL_GPU_TYPES, normalize_modal_gpu_selection
from vast_config_node import VastAILeaseConfiguration, profile_from_inputs
logger = logging.getLogger(__name__)
REMOTE_CONFIGURATION_IO_TYPE = "REMOTE_CONFIGURATION"
REMOTE_CONFIGURATION_SET_IO_TYPE = "REMOTE_CONFIGURATION_SET"
REMOTE_EXECUTION_CONFIGURATOR_NODE_ID = "RemoteExecutionConfigurator"
MODAL_REMOTE_CONFIGURATION_NODE_ID = "ModalRemoteConfiguration"
VAST_REMOTE_CONFIGURATION_NODE_ID = "VastRemoteConfiguration"
SSH_REMOTE_CONFIGURATION_NODE_ID = "SshRemoteConfiguration"
R2_STORAGE_CONFIGURATION_NODE_ID = "R2StorageBackingConfiguration"
REMOTE_CONFIGURATION_INPUT_GROUP = "configurations"
REMOTE_CONFIGURATION_INPUT_PREFIX = "configuration_"
REMOTE_CONFIGURATION_MAX_INPUTS = 32
REMOTE_CONFIGURATION_NODE_IDS = frozenset(
{
REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
MODAL_REMOTE_CONFIGURATION_NODE_ID,
VAST_REMOTE_CONFIGURATION_NODE_ID,
SSH_REMOTE_CONFIGURATION_NODE_ID,
R2_STORAGE_CONFIGURATION_NODE_ID,
}
)
RemoteConfigurationType = io.Custom(REMOTE_CONFIGURATION_IO_TYPE)
RemoteConfigurationSetType = io.Custom(REMOTE_CONFIGURATION_SET_IO_TYPE)
class StorageBackingConfigurationNode(io.ComfyNode):
"""Base ComfyUI node for workflow-scoped shared storage configuration."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Expose one concrete storage provider's non-secret settings."""
raise NotImplementedError
class R2StorageConfiguration(StorageBackingConfigurationNode):
"""Declare a Cloudflare R2 backing cache authorized by the local controller."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Expose portable R2 metadata while keeping credentials out of workflows."""
return io.Schema(
node_id=R2_STORAGE_CONFIGURATION_NODE_ID,
display_name="R2 Storage Configuration",
category="Remote Execution/Configuration/Storage",
description=(
"Provide a shared Cloudflare R2 backing cache to Vast.ai and SSH "
"targets. Use Login on the node to authorize Cloudflare without "
"executing the workflow."
),
inputs=_r2_storage_inputs(),
outputs=[RemoteConfigurationType.Output()],
hidden=[io.Hidden.unique_id],
is_experimental=True,
)
@classmethod
def execute(cls, **inputs: Any) -> io.NodeOutput:
"""Build one credential-reference-only R2 storage configuration value."""
configuration_id = _hidden_unique_id(cls) or str(
inputs.pop("unique_id", "r2-storage")
)
return io.NodeOutput(
r2_storage_configuration_from_inputs(configuration_id, inputs)
)
def _r2_storage_inputs() -> list[io.Input]:
"""Return non-secret workflow widgets shared by every R2 node instance."""
return [
io.String.Input(
"configuration_name",
default="r2-cache",
tooltip="Workflow-local label for this storage backing.",
),
io.String.Input(
"account_id",
default="",
tooltip="Filled automatically after Cloudflare Login.",
),
io.String.Input(
"bucket",
default="comfy-modal-cache",
tooltip="Private R2 bucket created or reused during Login.",
),
io.Combo.Input(
"jurisdiction",
options=["default", "eu", "fedramp", "us"],
default="default",
tooltip="R2 data jurisdiction and S3 endpoint family.",
),
io.Combo.Input(
"write_back_mode",
options=["async", "off"],
default="async",
tooltip="Populate R2 during workflow-idle time, or use R2 read-only.",
),
io.String.Input(
"key_prefix",
default="comfy-modal-cache/v1/blobs/sha256",
advanced=True,
tooltip="Immutable content-addressed object namespace.",
),
io.String.Input(
"credential_id",
default="",
advanced=True,
tooltip=(
"Random local keyring reference generated by the UI; it is not a "
"credential."
),
),
]
class ModalConfiguration(io.ComfyNode):
"""Declare one workflow-scoped Modal capacity pool."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Expose a Modal GPU type and concurrent instance limit."""
return io.Schema(
node_id=MODAL_REMOTE_CONFIGURATION_NODE_ID,
display_name="Modal Configuration",
category="Remote Execution/Configuration",
description=(
"Contribute one Modal GPU type and capacity limit to a Remote "
"Execution Configurator."
),
inputs=[
io.String.Input(
"configuration_name",
default="modal-default",
tooltip="Unique workflow-local name for this Modal capacity pool.",
),
io.Combo.Input(
"gpu_type",
options=list(MODAL_GPU_TYPES),
default="RTX-PRO-6000",
tooltip="Modal GPU type used by this capacity pool.",
),
io.Int.Input(
"instance_count",
default=1,
min=1,
max=32,
step=1,
tooltip="Maximum concurrent Modal containers for this pool.",
),
],
outputs=[RemoteConfigurationType.Output()],
hidden=[io.Hidden.unique_id],
is_experimental=True,
)
@classmethod
def execute(cls, **inputs: Any) -> io.NodeOutput:
"""Build one validated Modal configuration value."""
configuration_id = _hidden_unique_id(cls) or str(
inputs.pop("unique_id", "modal-configuration")
)
return io.NodeOutput(modal_configuration_from_inputs(configuration_id, inputs))
class VastConfiguration(io.ComfyNode):
"""Declare one workflow-scoped Vast.ai marketplace capacity pool."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Reuse the established Vast.ai marketplace controls with a typed output."""
legacy_schema = VastAILeaseConfiguration.define_schema()
return io.Schema(
node_id=VAST_REMOTE_CONFIGURATION_NODE_ID,
display_name="Vast.ai Configuration",
category="Remote Execution/Configuration",
description=(
"Contribute one Vast.ai marketplace search and managed-instance "
"limit to a Remote Execution Configurator."
),
inputs=list(legacy_schema.inputs),
outputs=[RemoteConfigurationType.Output()],
hidden=[io.Hidden.unique_id],
is_experimental=True,
)
@classmethod
def execute(cls, **inputs: Any) -> io.NodeOutput:
"""Build one validated Vast.ai configuration value."""
configuration_id = _hidden_unique_id(cls) or str(
inputs.pop("unique_id", "vast-configuration")
)
return io.NodeOutput(vast_configuration_from_inputs(configuration_id, inputs))
class SshConfiguration(io.ComfyNode):
"""Declare one workflow-scoped SSH Docker execution host."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Expose portable, credential-free host and scheduling controls."""
return io.Schema(
node_id=SSH_REMOTE_CONFIGURATION_NODE_ID,
display_name="SSH Configuration",
category="Remote Execution/Configuration",
description=(
"Contribute one SSH-accessible Docker host to a Remote Execution "
"Configurator. Authentication remains in the local SSH agent/config."
),
inputs=[
io.String.Input(
"environment_id",
default="ssh-host",
tooltip="Stable lowercase workflow-local environment identity.",
),
io.String.Input(
"display_name",
default="SSH host",
tooltip="Human-readable name used in plans and errors.",
),
io.String.Input(
"ssh_target",
default="",
tooltip=(
"OpenSSH destination or alias that works "
"non-interactively."
),
),
io.String.Input(
"cost_usd_per_hour",
default="Unknown",
tooltip=(
"Marginal hourly cost used by cost-aware planning. Use "
"Unknown when no trustworthy price is available."
),
),
io.Int.Input(
"maximum_workers",
default=1,
min=1,
max=32,
step=1,
tooltip="Maximum worker containers addressed on this host.",
),
io.Float.Input(
"reserve_vram_gb",
default=0.0,
min=0.0,
max=4096.0,
step=0.1,
advanced=True,
tooltip="GPU VRAM withheld from scheduler admission checks.",
),
io.String.Input(
"tags",
default="",
optional=True,
advanced=True,
tooltip="Comma-separated placement tags.",
),
io.String.Input(
"docker_env_file",
default="",
optional=True,
advanced=True,
tooltip=(
"Optional absolute env-file path already present on the host."
),
),
],
outputs=[RemoteConfigurationType.Output()],
hidden=[io.Hidden.unique_id],
is_experimental=True,
)
@classmethod
def execute(cls, **inputs: Any) -> io.NodeOutput:
"""Build one validated SSH host configuration value."""
configuration_id = _hidden_unique_id(cls) or str(
inputs.pop("unique_id", inputs.get("environment_id", "ssh-host"))
)
return io.NodeOutput(ssh_configuration_from_inputs(configuration_id, inputs))
class RemoteExecutionConfigurator(io.ComfyNode):
"""Anchor and validate the complete remote capacity declaration."""
@classmethod
def define_schema(cls) -> io.Schema:
"""Accept a variable number of typed remote configurations."""
configuration_template = io.Autogrow.TemplatePrefix(
input=RemoteConfigurationType.Input("configuration"),
prefix=REMOTE_CONFIGURATION_INPUT_PREFIX,
min=1,
max=REMOTE_CONFIGURATION_MAX_INPUTS,
)
return io.Schema(
node_id=REMOTE_EXECUTION_CONFIGURATOR_NODE_ID,
display_name="Remote Execution Configurator",
category="Remote Execution/Configuration",
description=(
"Collect the Modal, Vast.ai, and SSH capacity pools plus shared "
"storage backings available to this workflow's remote execution "
"planner."
),
inputs=[
io.Autogrow.Input(
REMOTE_CONFIGURATION_INPUT_GROUP,
template=configuration_template,
)
],
outputs=[RemoteConfigurationSetType.Output()],
is_output_node=True,
is_experimental=True,
)
@classmethod
def execute(cls, configurations: io.Autogrow.Type) -> io.NodeOutput:
"""Assemble ordered configuration values into one validated set."""
return io.NodeOutput(configuration_set_from_values(configurations))
def modal_configuration_from_inputs(
configuration_id: str,
inputs: Mapping[str, Any],
) -> ModalRemoteConfiguration:
"""Build one Modal configuration from queued widget inputs."""
return ModalRemoteConfiguration(
configuration_id=str(configuration_id).strip(),
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)),
)
def vast_configuration_from_inputs(
configuration_id: str,
inputs: Mapping[str, Any],
) -> VastRemoteConfiguration:
"""Build one Vast configuration from queued widget inputs."""
profile = profile_from_inputs(str(configuration_id).strip(), inputs)
return VastRemoteConfiguration(
configuration_id=profile.profile_id,
display_name=profile.profile_name,
profile=profile,
)
def ssh_configuration_from_inputs(
configuration_id: str,
inputs: Mapping[str, Any],
) -> SshRemoteConfiguration:
"""Build one SSH configuration from queued widget inputs."""
environment_id = str(inputs.get("environment_id") or configuration_id).strip()
raw_tags = str(inputs.get("tags") or "")
host = SshHostConfig(
environment_id=environment_id,
display_name=str(inputs.get("display_name") or environment_id).strip(),
ssh_target=str(inputs.get("ssh_target") or "").strip(),
cost_usd_per_second=_optional_hourly_cost_per_second(
inputs.get("cost_usd_per_hour")
),
maximum_workers=int(inputs.get("maximum_workers", 1)),
reserve_vram_bytes=round(float(inputs.get("reserve_vram_gb", 0.0)) * 1024**3),
tags=frozenset(tag.strip() for tag in raw_tags.split(",") if tag.strip()),
docker_env_file=(
str(inputs["docker_env_file"]).strip()
if inputs.get("docker_env_file")
else None
),
)
return SshRemoteConfiguration(
configuration_id=host.environment_id,
display_name=host.display_name,
host=host,
)
def r2_storage_configuration_from_inputs(
configuration_id: str,
inputs: Mapping[str, Any],
) -> R2StorageBackingConfiguration:
"""Build one non-secret R2 storage configuration from queued widget inputs."""
normalized_id = str(configuration_id).strip()
return R2StorageBackingConfiguration(
configuration_id=normalized_id,
display_name=str(inputs.get("configuration_name") or "r2-cache").strip(),
account_id=str(inputs.get("account_id") or "").strip(),
bucket=str(inputs.get("bucket") or "").strip(),
credential_id=str(inputs.get("credential_id") or normalized_id).strip(),
jurisdiction=str(inputs.get("jurisdiction") or "default").strip().casefold(),
key_prefix=str(
inputs.get("key_prefix") or "comfy-modal-cache/v1/blobs/sha256"
).strip(),
write_back_mode=str(inputs.get("write_back_mode") or "async")
.strip()
.casefold(),
)
def configuration_set_from_values(
values: Mapping[str, Any],
) -> RemoteConfigurationSet:
"""Validate ordered autogrow values as remote configurations."""
configurations: list[RemoteExecutionConfiguration] = []
for input_name, value in sorted(
values.items(),
key=lambda item: _configuration_input_sort_key(str(item[0])),
):
if not isinstance(value, RemoteExecutionConfiguration):
raise TypeError(
f"Remote configurator input {input_name!r} did not receive a "
"RemoteExecutionConfiguration value."
)
configurations.append(value)
return RemoteConfigurationSet(tuple(configurations))
def compile_remote_configuration_set(
prompt: Mapping[str, Any],
) -> RemoteConfigurationSet | None:
"""Compile the connected configurator branch before ComfyUI node execution."""
configurator_nodes = [
(str(node_id), prompt_node)
for node_id, prompt_node in prompt.items()
if isinstance(prompt_node, Mapping)
and str(prompt_node.get("class_type") or "")
== REMOTE_EXECUTION_CONFIGURATOR_NODE_ID
]
if not configurator_nodes:
return None
if len(configurator_nodes) > 1:
raise ValueError(
"A workflow may contain only one Remote Execution Configurator."
)
configurator_id, configurator_node = configurator_nodes[0]
raw_inputs = configurator_node.get("inputs")
if not isinstance(raw_inputs, Mapping):
raise ValueError(
f"Remote Execution Configurator {configurator_id!r} has invalid inputs."
)
connected_inputs = _serialized_configuration_inputs(raw_inputs)
if not connected_inputs:
raise ValueError(
"Remote Execution Configurator requires at least one connected "
"REMOTE_CONFIGURATION input."
)
configurations: list[RemoteExecutionConfiguration] = []
linked_node_ids: set[str] = set()
for input_name, input_value in connected_inputs:
source_node_id = _configuration_link_source(input_name, input_value)
if source_node_id in linked_node_ids:
raise ValueError(
f"Remote configuration node {source_node_id!r} is connected "
"more than once."
)
linked_node_ids.add(source_node_id)
source_node = prompt.get(source_node_id)
if not isinstance(source_node, Mapping):
raise ValueError(
f"Remote configurator input {input_name!r} references missing node "
f"{source_node_id!r}."
)
source_inputs = source_node.get("inputs")
if not isinstance(source_inputs, Mapping):
raise ValueError(
f"Remote configuration node {source_node_id!r} has invalid inputs."
)
configurations.append(
_configuration_from_prompt_node(
source_node_id,
str(source_node.get("class_type") or ""),
source_inputs,
)
)
return RemoteConfigurationSet(tuple(configurations))
def _serialized_configuration_inputs(
raw_inputs: Mapping[str, Any],
) -> list[tuple[str, Any]]:
"""Return Autogrow links from serialized or reconstructed input shapes."""
flattened_prefix = (
f"{REMOTE_CONFIGURATION_INPUT_GROUP}."
f"{REMOTE_CONFIGURATION_INPUT_PREFIX}"
)
candidates: list[tuple[str, str, Any]] = []
for raw_name, value in raw_inputs.items():
input_name = str(raw_name)
if input_name.startswith(flattened_prefix):
leaf_name = input_name.removeprefix(
f"{REMOTE_CONFIGURATION_INPUT_GROUP}."
)
candidates.append((leaf_name, input_name, value))
continue
if input_name.startswith(REMOTE_CONFIGURATION_INPUT_PREFIX):
candidates.append((input_name, input_name, value))
continue
if (
input_name == REMOTE_CONFIGURATION_INPUT_GROUP
and isinstance(value, Mapping)
):
for nested_name, nested_value in value.items():
leaf_name = str(nested_name)
if not leaf_name.startswith(REMOTE_CONFIGURATION_INPUT_PREFIX):
continue
candidates.append(
(
leaf_name,
f"{REMOTE_CONFIGURATION_INPUT_GROUP}.{leaf_name}",
nested_value,
)
)
leaf_names = [leaf_name for leaf_name, _, _ in candidates]
if len(leaf_names) != len(set(leaf_names)):
raise ValueError(
"Remote Execution Configurator contains duplicate Autogrow input slots."
)
return [
(input_name, value)
for _, input_name, value in sorted(
candidates,
key=lambda item: _configuration_input_sort_key(item[0]),
)
]
def _configuration_link_source(input_name: str, value: Any) -> str:
"""Return the source node ID from one typed configurator link."""
if (
not isinstance(value, list)
or len(value) != 2
or isinstance(value[1], bool)
or int(value[1]) != 0
):
raise ValueError(
f"Remote configurator input {input_name!r} must connect to output 0 "
"of a configuration node."
)
return str(value[0])
def _configuration_from_prompt_node(
node_id: str,
class_type: str,
inputs: Mapping[str, Any],
) -> RemoteExecutionConfiguration:
"""Build the typed configuration represented by one queued prompt node."""
if class_type == MODAL_REMOTE_CONFIGURATION_NODE_ID:
return modal_configuration_from_inputs(node_id, inputs)
if class_type == VAST_REMOTE_CONFIGURATION_NODE_ID:
return vast_configuration_from_inputs(node_id, inputs)
if class_type == SSH_REMOTE_CONFIGURATION_NODE_ID:
return ssh_configuration_from_inputs(node_id, inputs)
if class_type == R2_STORAGE_CONFIGURATION_NODE_ID:
return r2_storage_configuration_from_inputs(node_id, inputs)
raise ValueError(
f"Remote configurator source node {node_id!r} has unsupported type "
f"{class_type!r}."
)
def _configuration_input_sort_key(input_name: str) -> tuple[int, str]:
"""Return numeric autogrow ordering with a stable fallback."""
suffix = input_name.removeprefix(REMOTE_CONFIGURATION_INPUT_PREFIX)
try:
return (int(suffix), input_name)
except ValueError:
return (REMOTE_CONFIGURATION_MAX_INPUTS + 1, input_name)
def _optional_hourly_cost_per_second(value: Any) -> float | None:
"""Convert an optional hourly workflow price to a per-second scheduler rate."""
normalized = str(value if value is not None else "").strip()
if not normalized or normalized.casefold() in {"unknown", "any"}:
return None
try:
hourly_cost = float(normalized)
except ValueError as exc:
raise ValueError(
"cost_usd_per_hour must be Unknown or a non-negative number."
) from exc
if hourly_cost < 0:
raise ValueError("cost_usd_per_hour must not be negative.")
return hourly_cost / 3600.0
def _hidden_unique_id(node_class: type[io.ComfyNode]) -> str | None:
"""Return the current v3 hidden unique ID when ComfyUI supplied one."""
hidden = getattr(node_class, "hidden", None)
unique_id = getattr(hidden, "unique_id", None)
if unique_id is None:
return None
return str(unique_id)
__all__ = [
"MODAL_REMOTE_CONFIGURATION_NODE_ID",
"ModalConfiguration",
"REMOTE_CONFIGURATION_IO_TYPE",
"REMOTE_CONFIGURATION_INPUT_GROUP",
"REMOTE_CONFIGURATION_NODE_IDS",
"REMOTE_CONFIGURATION_SET_IO_TYPE",
"REMOTE_EXECUTION_CONFIGURATOR_NODE_ID",
"RemoteExecutionConfigurator",
"R2_STORAGE_CONFIGURATION_NODE_ID",
"R2StorageConfiguration",
"SSH_REMOTE_CONFIGURATION_NODE_ID",
"SshConfiguration",
"StorageBackingConfigurationNode",
"VAST_REMOTE_CONFIGURATION_NODE_ID",
"VastConfiguration",
"compile_remote_configuration_set",
"configuration_set_from_values",
"modal_configuration_from_inputs",
"r2_storage_configuration_from_inputs",
"ssh_configuration_from_inputs",
"vast_configuration_from_inputs",
]