556 lines
20 KiB
Python
556 lines
20 KiB
Python
"""Disconnected ComfyUI v3 node for workflow-scoped Vast.ai lease profiles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from typing import Any, Mapping
|
|
|
|
from comfy_api.latest import _io as io
|
|
|
|
if __package__:
|
|
from .vast_models import (
|
|
VAST_CONFIG_NODE_ID,
|
|
VAST_DEFAULT_EXCLUDED_COUNTRY_CODES,
|
|
VAST_DEFAULT_IDLE_RETENTION_HOURS,
|
|
VastResourceProfile,
|
|
)
|
|
else: # pragma: no cover - direct ComfyUI loading fallback.
|
|
from vast_models import (
|
|
VAST_CONFIG_NODE_ID,
|
|
VAST_DEFAULT_EXCLUDED_COUNTRY_CODES,
|
|
VAST_DEFAULT_IDLE_RETENTION_HOURS,
|
|
VastResourceProfile,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
VAST_ANY_VALUE = "Any"
|
|
_VAST_VERIFIED_ONLY_VALUE = "Verified only"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastSelection:
|
|
"""Describe one distinct GPU selection rendered by the configuration node."""
|
|
|
|
gpu_name: str
|
|
gpu_count: int
|
|
gpu_ram_mb: int
|
|
hourly_cost_usd: float
|
|
instance_count: int
|
|
component_count: int
|
|
|
|
|
|
class VastAILeaseConfiguration(io.ComfyNode):
|
|
"""Declare one disconnected Vast.ai resource and retention profile."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Expose portable non-secret capacity and spending controls."""
|
|
return io.Schema(
|
|
node_id=VAST_CONFIG_NODE_ID,
|
|
display_name="Vast.ai Lease Configuration",
|
|
category="Remote Execution/Vast.ai",
|
|
description=(
|
|
"Configure a Vast.ai capacity pool for this workflow. This node "
|
|
"remains disconnected; remote placement detects it automatically "
|
|
"at queue time."
|
|
),
|
|
inputs=[
|
|
io.String.Input(
|
|
"profile_name",
|
|
default="vast-default",
|
|
tooltip="Unique workflow-local name for this Vast capacity pool.",
|
|
),
|
|
io.String.Input(
|
|
"gpu_count",
|
|
default=VAST_ANY_VALUE,
|
|
tooltip="Use Any or enter the exact number of GPUs to rent.",
|
|
),
|
|
io.String.Input(
|
|
"minimum_gpu_vram_gb",
|
|
default=VAST_ANY_VALUE,
|
|
tooltip=(
|
|
"Use Any to rely on the workflow memory estimate, or enter the "
|
|
"minimum VRAM in GiB for each GPU."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_total_tflops",
|
|
default=VAST_ANY_VALUE,
|
|
tooltip=(
|
|
"Use Any or enter minimum theoretical aggregate TFLOPS. Vast's "
|
|
"cross-architecture TFLOPS values are not directly comparable."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_cpu_ram_gb",
|
|
default=VAST_ANY_VALUE,
|
|
tooltip=(
|
|
"Use Any to rely on the workflow memory estimate, or enter the "
|
|
"minimum system RAM in GiB."
|
|
),
|
|
),
|
|
io.Float.Input(
|
|
"allocated_disk_gb",
|
|
default=200.0,
|
|
min=8.0,
|
|
max=10000.0,
|
|
step=1.0,
|
|
tooltip=(
|
|
"Instance disk allocation. Vast does not allow resizing it "
|
|
"after creation."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"maximum_hourly_cost_usd",
|
|
default=VAST_ANY_VALUE,
|
|
tooltip=(
|
|
"Use Any for no hard hourly price ceiling, or enter a positive "
|
|
"USD limit. Any may permit an expensive Vast-only rental."
|
|
),
|
|
),
|
|
io.Float.Input(
|
|
"idle_retention_hours",
|
|
default=VAST_DEFAULT_IDLE_RETENTION_HOURS,
|
|
min=0.0,
|
|
max=24.0 * 365.0,
|
|
step=1.0,
|
|
tooltip=(
|
|
"Keep the running instance after its last activity for this "
|
|
"many hours, then destroy it. Vast bills during this period."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_cpu_cores",
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Use Any or enter the minimum effective CPU core allocation."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_dlperf",
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip="Use Any or enter an optional Vast DLPerf floor.",
|
|
),
|
|
io.String.Input(
|
|
"minimum_download_mb_per_second",
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Use Any or enter the minimum advertised internet download "
|
|
"rate in MB/s."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_reliability",
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Use Any or enter the minimum Vast host reliability score "
|
|
"between 0 and 1."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"minimum_offer_duration_days",
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Use Any or require the offer to remain available for at least "
|
|
"this many days. This is separate from idle retention."
|
|
),
|
|
),
|
|
io.Combo.Input(
|
|
"verified_hosts_only",
|
|
options=[VAST_ANY_VALUE, _VAST_VERIFIED_ONLY_VALUE],
|
|
default=VAST_ANY_VALUE,
|
|
advanced=True,
|
|
tooltip="Use Any or exclude unverified Vast marketplace hosts.",
|
|
),
|
|
io.String.Input(
|
|
"allowed_geolocations",
|
|
default=VAST_ANY_VALUE,
|
|
optional=True,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Use Any or enter comma-separated Vast locations or country "
|
|
"codes, such as US, CA."
|
|
),
|
|
),
|
|
io.String.Input(
|
|
"excluded_country_codes",
|
|
default=", ".join(VAST_DEFAULT_EXCLUDED_COUNTRY_CODES),
|
|
optional=True,
|
|
advanced=True,
|
|
tooltip=(
|
|
"Comma-separated two-letter country codes to exclude from "
|
|
"Vast marketplace searches, such as CN, RU. Leave blank to "
|
|
"disable country exclusions."
|
|
),
|
|
),
|
|
io.Int.Input(
|
|
"maximum_instances",
|
|
default=1,
|
|
min=1,
|
|
max=16,
|
|
step=1,
|
|
advanced=True,
|
|
tooltip="Maximum simultaneously managed instances for this pool.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.String.Output(
|
|
display_name="STRING",
|
|
tooltip=(
|
|
"Markdown describing the Vast.ai GPU type or types selected "
|
|
"for this profile at queue time."
|
|
),
|
|
)
|
|
],
|
|
hidden=[io.Hidden.unique_id, io.Hidden.prompt],
|
|
is_output_node=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **inputs: Any) -> io.NodeOutput:
|
|
"""Validate the profile and describe its actual queue-time Vast selection."""
|
|
unique_id = _hidden_unique_id(cls) or str(
|
|
inputs.pop("unique_id", "vast-config")
|
|
)
|
|
profile = profile_from_inputs(unique_id, inputs)
|
|
return io.NodeOutput(
|
|
vast_selection_markdown(
|
|
_hidden_prompt(cls),
|
|
profile_id=profile.profile_id,
|
|
profile_name=profile.profile_name,
|
|
)
|
|
)
|
|
|
|
@classmethod
|
|
def fingerprint_inputs(cls, **inputs: Any) -> str:
|
|
"""Invalidate cached Markdown whenever queue-time placement changes."""
|
|
unique_id = _hidden_unique_id(cls) or str(
|
|
inputs.get("unique_id", "vast-config")
|
|
)
|
|
profile = profile_from_inputs(unique_id, inputs)
|
|
return vast_selection_markdown(
|
|
_hidden_prompt(cls),
|
|
profile_id=profile.profile_id,
|
|
profile_name=profile.profile_name,
|
|
)
|
|
|
|
|
|
def profile_from_inputs(
|
|
profile_id: str,
|
|
inputs: Mapping[str, Any],
|
|
) -> VastResourceProfile:
|
|
"""Build one validated Vast resource profile from queued node inputs."""
|
|
locations = tuple(
|
|
location.strip()
|
|
for location in str(inputs.get("allowed_geolocations") or "").split(",")
|
|
if location.strip() and not _is_any(location)
|
|
)
|
|
excluded_country_codes = _excluded_country_codes(
|
|
inputs.get(
|
|
"excluded_country_codes",
|
|
", ".join(VAST_DEFAULT_EXCLUDED_COUNTRY_CODES),
|
|
)
|
|
)
|
|
return VastResourceProfile(
|
|
profile_id=str(profile_id).strip(),
|
|
profile_name=str(inputs.get("profile_name") or "vast-default").strip(),
|
|
gpu_count=_optional_positive_int(inputs.get("gpu_count"), "gpu_count"),
|
|
minimum_gpu_ram_mb=_optional_gb_to_mb(
|
|
inputs.get("minimum_gpu_vram_gb"),
|
|
"minimum_gpu_vram_gb",
|
|
),
|
|
minimum_total_flops=_optional_non_negative_float(
|
|
inputs.get("minimum_total_tflops"),
|
|
"minimum_total_tflops",
|
|
),
|
|
minimum_cpu_ram_mb=_optional_gb_to_mb(
|
|
inputs.get("minimum_cpu_ram_gb"),
|
|
"minimum_cpu_ram_gb",
|
|
),
|
|
minimum_cpu_cores=_optional_non_negative_float(
|
|
inputs.get("minimum_cpu_cores"),
|
|
"minimum_cpu_cores",
|
|
),
|
|
allocated_disk_gb=float(inputs.get("allocated_disk_gb", 200.0)),
|
|
maximum_hourly_cost_usd=_optional_positive_float(
|
|
inputs.get("maximum_hourly_cost_usd"),
|
|
"maximum_hourly_cost_usd",
|
|
),
|
|
idle_retention_seconds=float(
|
|
inputs.get(
|
|
"idle_retention_hours", VAST_DEFAULT_IDLE_RETENTION_HOURS
|
|
)
|
|
)
|
|
* 3600.0,
|
|
minimum_offer_duration_seconds=_optional_days_to_seconds(
|
|
inputs.get("minimum_offer_duration_days")
|
|
),
|
|
minimum_reliability=_optional_bounded_float(
|
|
inputs.get("minimum_reliability"),
|
|
"minimum_reliability",
|
|
maximum=1.0,
|
|
),
|
|
minimum_dlperf=_optional_non_negative_float(
|
|
inputs.get("minimum_dlperf"),
|
|
"minimum_dlperf",
|
|
),
|
|
minimum_download_mb_per_second=_optional_non_negative_float(
|
|
inputs.get("minimum_download_mb_per_second"),
|
|
"minimum_download_mb_per_second",
|
|
),
|
|
verified_only=_verified_only(inputs.get("verified_hosts_only")),
|
|
allowed_geolocations=locations,
|
|
excluded_country_codes=excluded_country_codes,
|
|
maximum_instances=int(inputs.get("maximum_instances", 1)),
|
|
)
|
|
|
|
|
|
def _excluded_country_codes(value: Any) -> tuple[str, ...]:
|
|
"""Normalize a comma-separated list of unique ISO-style country codes."""
|
|
country_codes: list[str] = []
|
|
for raw_country_code in str(value or "").split(","):
|
|
country_code = raw_country_code.strip().upper()
|
|
if not country_code or _is_any(country_code):
|
|
continue
|
|
if country_code not in country_codes:
|
|
country_codes.append(country_code)
|
|
return tuple(country_codes)
|
|
|
|
|
|
def extract_vast_profiles(prompt: Mapping[str, Any]) -> tuple[VastResourceProfile, ...]:
|
|
"""Return every disconnected Vast configuration in one executable prompt."""
|
|
profiles: list[VastResourceProfile] = []
|
|
names: set[str] = set()
|
|
for node_id, prompt_node in prompt.items():
|
|
if not isinstance(prompt_node, Mapping):
|
|
continue
|
|
if str(prompt_node.get("class_type") or "") != VAST_CONFIG_NODE_ID:
|
|
continue
|
|
raw_inputs = prompt_node.get("inputs")
|
|
if not isinstance(raw_inputs, Mapping):
|
|
raise ValueError(f"Vast configuration node {node_id!r} has invalid inputs.")
|
|
profile = profile_from_inputs(str(node_id), raw_inputs)
|
|
normalized_name = profile.profile_name.casefold()
|
|
if normalized_name in names:
|
|
raise ValueError(
|
|
f"Vast profile name {profile.profile_name!r} appears more than once."
|
|
)
|
|
names.add(normalized_name)
|
|
profiles.append(profile)
|
|
return tuple(sorted(profiles, key=lambda profile: profile.profile_id))
|
|
|
|
|
|
def vast_selection_markdown(
|
|
prompt: Mapping[str, Any] | None,
|
|
*,
|
|
profile_id: str,
|
|
profile_name: str,
|
|
) -> str:
|
|
"""Return Markdown describing distinct Vast GPU selections for one profile."""
|
|
selections = _vast_selections(prompt, profile_id=profile_id)
|
|
heading = f"## Vast.ai selection for `{_escape_markdown_code(profile_name)}`"
|
|
if not selections:
|
|
return (
|
|
f"{heading}\n\n"
|
|
"No Vast.ai node type was selected for this execution."
|
|
)
|
|
rows = [
|
|
"| GPU type | GPUs per instance | VRAM per GPU | Hourly price | "
|
|
"Instances | Components |",
|
|
"|---|---:|---:|---:|---:|---:|",
|
|
]
|
|
rows.extend(
|
|
(
|
|
"| {gpu} | {gpu_count} | {vram:.3f} GiB | ${price:.3f} | "
|
|
"{instances} | {components} |"
|
|
).format(
|
|
gpu=_escape_markdown_cell(selection.gpu_name),
|
|
gpu_count=selection.gpu_count,
|
|
vram=selection.gpu_ram_mb / 1024.0,
|
|
price=selection.hourly_cost_usd,
|
|
instances=selection.instance_count,
|
|
components=selection.component_count,
|
|
)
|
|
for selection in selections
|
|
)
|
|
return f"{heading}\n\n" + "\n".join(rows)
|
|
|
|
|
|
def _vast_selections(
|
|
prompt: Mapping[str, Any] | None,
|
|
*,
|
|
profile_id: str,
|
|
) -> tuple[VastSelection, ...]:
|
|
"""Extract and aggregate safe Vast placement metadata from a rewritten prompt."""
|
|
if not isinstance(prompt, Mapping):
|
|
return ()
|
|
components_by_selection: dict[
|
|
tuple[str, int, int, float], set[str]
|
|
] = defaultdict(set)
|
|
instances_by_selection: dict[
|
|
tuple[str, int, int, float], set[int]
|
|
] = defaultdict(set)
|
|
for node_id, prompt_node in prompt.items():
|
|
if not isinstance(prompt_node, Mapping):
|
|
continue
|
|
inputs = prompt_node.get("inputs")
|
|
if not isinstance(inputs, Mapping):
|
|
continue
|
|
payload = inputs.get("original_node_data")
|
|
if not isinstance(payload, Mapping):
|
|
continue
|
|
if str(payload.get("execution_provider") or "") != "vast":
|
|
continue
|
|
if str(payload.get("vast_profile_id") or "") != profile_id:
|
|
continue
|
|
key = (
|
|
str(payload.get("vast_gpu_name") or "Unknown GPU"),
|
|
int(payload.get("vast_gpu_count") or 1),
|
|
int(payload.get("vast_gpu_ram_mb") or 0),
|
|
float(payload.get("vast_hourly_cost_usd") or 0.0),
|
|
)
|
|
components_by_selection[key].add(str(payload.get("component_id") or node_id))
|
|
instance_id = int(payload.get("vast_instance_id") or 0)
|
|
if instance_id > 0:
|
|
instances_by_selection[key].add(instance_id)
|
|
return tuple(
|
|
VastSelection(
|
|
gpu_name=key[0],
|
|
gpu_count=key[1],
|
|
gpu_ram_mb=key[2],
|
|
hourly_cost_usd=key[3],
|
|
instance_count=max(1, len(instances_by_selection[key])),
|
|
component_count=len(component_ids),
|
|
)
|
|
for key, component_ids in sorted(
|
|
components_by_selection.items(),
|
|
key=lambda item: (item[0][3], item[0][0]),
|
|
)
|
|
)
|
|
|
|
|
|
def _hidden_prompt(
|
|
node_class: type[VastAILeaseConfiguration],
|
|
) -> Mapping[str, Any] | None:
|
|
"""Return the v3 hidden prompt when ComfyUI supplied one."""
|
|
hidden = getattr(node_class, "hidden", None)
|
|
prompt = getattr(hidden, "prompt", None)
|
|
return prompt if isinstance(prompt, Mapping) else None
|
|
|
|
|
|
def _hidden_unique_id(node_class: type[VastAILeaseConfiguration]) -> str | None:
|
|
"""Return the v3 hidden node identity when ComfyUI supplied one."""
|
|
hidden = getattr(node_class, "hidden", None)
|
|
unique_id = str(getattr(hidden, "unique_id", "") or "").strip()
|
|
return unique_id or None
|
|
|
|
|
|
def _is_any(value: Any) -> bool:
|
|
"""Return whether a queued selector explicitly requests no constraint."""
|
|
return value is None or str(value).strip().casefold() in {"", "any"}
|
|
|
|
|
|
def _optional_positive_int(value: Any, field_name: str) -> int | None:
|
|
"""Parse Any or a positive integer selector."""
|
|
if _is_any(value):
|
|
return None
|
|
if isinstance(value, bool):
|
|
raise ValueError(f"{field_name} must be Any or a positive integer.")
|
|
try:
|
|
normalized = int(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{field_name} must be Any or a positive integer.") from exc
|
|
if normalized <= 0 or float(value) != normalized:
|
|
raise ValueError(f"{field_name} must be Any or a positive integer.")
|
|
return normalized
|
|
|
|
|
|
def _optional_gb_to_mb(value: Any, field_name: str) -> int | None:
|
|
"""Convert Any or a GiB widget value to Vast's integer memory unit."""
|
|
normalized = _optional_non_negative_float(value, field_name)
|
|
return None if normalized is None else int(round(normalized * 1024.0))
|
|
|
|
|
|
def _optional_non_negative_float(value: Any, field_name: str) -> float | None:
|
|
"""Parse Any or a finite non-negative floating-point selector."""
|
|
return _optional_bounded_float(value, field_name, minimum=0.0)
|
|
|
|
|
|
def _optional_positive_float(value: Any, field_name: str) -> float | None:
|
|
"""Parse Any or a finite positive floating-point selector."""
|
|
normalized = _optional_bounded_float(value, field_name, minimum=0.0)
|
|
if normalized is not None and normalized <= 0:
|
|
raise ValueError(f"{field_name} must be Any or a positive number.")
|
|
return normalized
|
|
|
|
|
|
def _optional_bounded_float(
|
|
value: Any,
|
|
field_name: str,
|
|
*,
|
|
minimum: float = 0.0,
|
|
maximum: float | None = None,
|
|
) -> float | None:
|
|
"""Parse Any or one finite floating-point selector within inclusive bounds."""
|
|
if _is_any(value):
|
|
return None
|
|
if isinstance(value, bool):
|
|
raise ValueError(f"{field_name} must be Any or a number.")
|
|
try:
|
|
normalized = float(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{field_name} must be Any or a number.") from exc
|
|
if not math.isfinite(normalized) or normalized < minimum:
|
|
raise ValueError(f"{field_name} must be Any or at least {minimum}.")
|
|
if maximum is not None and normalized > maximum:
|
|
raise ValueError(f"{field_name} must be Any or at most {maximum}.")
|
|
return normalized
|
|
|
|
|
|
def _optional_days_to_seconds(value: Any) -> float | None:
|
|
"""Convert Any or a non-negative day count to seconds."""
|
|
normalized = _optional_non_negative_float(
|
|
value,
|
|
"minimum_offer_duration_days",
|
|
)
|
|
return None if normalized is None else normalized * 86400.0
|
|
|
|
|
|
def _verified_only(value: Any) -> bool:
|
|
"""Parse the Any/Verified-only selector and legacy Boolean workflow values."""
|
|
if isinstance(value, bool):
|
|
return value
|
|
if _is_any(value):
|
|
return False
|
|
if str(value).strip().casefold() == _VAST_VERIFIED_ONLY_VALUE.casefold():
|
|
return True
|
|
raise ValueError("verified_hosts_only must be Any or Verified only.")
|
|
|
|
|
|
def _escape_markdown_cell(value: str) -> str:
|
|
"""Escape one untrusted marketplace label for a Markdown table cell."""
|
|
return value.replace("\\", "\\\\").replace("|", "\\|").replace("\n", " ")
|
|
|
|
|
|
def _escape_markdown_code(value: str) -> str:
|
|
"""Escape backticks in one short Markdown inline-code value."""
|
|
return value.replace("`", "\\`")
|
|
|
|
|
|
__all__ = [
|
|
"VastAILeaseConfiguration",
|
|
"VastSelection",
|
|
"extract_vast_profiles",
|
|
"profile_from_inputs",
|
|
"vast_selection_markdown",
|
|
]
|