Files
ttulttul-ComfyUI-Modal/execution_environments.py

636 lines
25 KiB
Python

"""Provider-neutral remote execution environment models and scheduling."""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Mapping, Protocol, Sequence
WORKFLOW_REMOTE_CONFIG_KEY = "remote_execution"
WORKFLOW_REMOTE_POLICY_KEY = "policy"
WORKFLOW_REMOTE_AUTO_PLACE_KEY = "auto_place"
WORKFLOW_REMOTE_PREFERRED_ENVIRONMENTS_KEY = "preferred_environment_ids"
WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY = "minimum_vram_gb"
class ExecutionProvider(str, Enum):
"""Identify one implementation of remote workflow execution."""
MODAL = "modal"
SSH_DOCKER = "ssh_docker"
VAST = "vast"
class ExecutionPolicy(str, Enum):
"""Describe how eligible workflow components select a provider."""
MODAL = "modal"
SELF_HOSTED = "self_hosted"
VAST = "vast"
AUTOMATIC = "automatic"
class EnvironmentHealth(str, Enum):
"""Describe whether a remote environment may receive new work."""
UNKNOWN = "unknown"
READY = "ready"
DEGRADED = "degraded"
UNAVAILABLE = "unavailable"
DRAINING = "draining"
@dataclass(frozen=True)
class GpuCapability:
"""Describe one GPU visible to a remote execution environment."""
uuid: str
name: str
total_vram_bytes: int
free_vram_bytes: int | None = None
compute_capability: str | None = None
driver_version: str | None = None
def to_dict(self) -> dict[str, Any]:
"""Return a JSON-compatible representation of this GPU."""
return {
"uuid": self.uuid,
"name": self.name,
"total_vram_bytes": self.total_vram_bytes,
"free_vram_bytes": self.free_vram_bytes,
"compute_capability": self.compute_capability,
"driver_version": self.driver_version,
}
@classmethod
def from_dict(cls, payload: Mapping[str, Any]) -> "GpuCapability":
"""Build one validated GPU capability from persisted data."""
return cls(
uuid=str(payload.get("uuid") or "").strip(),
name=str(payload.get("name") or "Unknown GPU").strip(),
total_vram_bytes=_non_negative_int(
payload.get("total_vram_bytes"), "total_vram_bytes"
),
free_vram_bytes=_optional_non_negative_int(
payload.get("free_vram_bytes"), "free_vram_bytes"
),
compute_capability=_optional_string(payload.get("compute_capability")),
driver_version=_optional_string(payload.get("driver_version")),
)
@dataclass(frozen=True)
class EnvironmentCapabilities:
"""Describe resources and runtime features discovered on one environment."""
architecture: str
operating_system: str
cpu_count: int
total_ram_bytes: int
available_ram_bytes: int | None
available_disk_bytes: int | None
docker_version: str
docker_rootless: bool
nvidia_container_runtime: bool
gpus: tuple[GpuCapability, ...] = ()
probed_at_epoch: float | None = None
reported_cost_usd_per_second: float | None = None
@property
def maximum_vram_bytes(self) -> int:
"""Return the largest total VRAM value exposed by one GPU."""
return max((gpu.total_vram_bytes for gpu in self.gpus), default=0)
@property
def maximum_free_vram_bytes(self) -> int:
"""Return the largest currently free VRAM value when reported."""
return max(
(
gpu.free_vram_bytes
for gpu in self.gpus
if gpu.free_vram_bytes is not None
),
default=0,
)
@property
def maximum_available_vram_bytes(self) -> int:
"""Return probed free VRAM when known, otherwise nameplate capacity."""
reported_free = [
gpu.free_vram_bytes for gpu in self.gpus if gpu.free_vram_bytes is not None
]
if reported_free:
return max(reported_free)
return self.maximum_vram_bytes
@property
def schedulable_ram_bytes(self) -> int:
"""Return currently available RAM when known, otherwise total RAM."""
if self.available_ram_bytes is not None:
return self.available_ram_bytes
return self.total_ram_bytes
def to_dict(self) -> dict[str, Any]:
"""Return a JSON-compatible representation of these capabilities."""
return {
"architecture": self.architecture,
"operating_system": self.operating_system,
"cpu_count": self.cpu_count,
"total_ram_bytes": self.total_ram_bytes,
"available_ram_bytes": self.available_ram_bytes,
"available_disk_bytes": self.available_disk_bytes,
"docker_version": self.docker_version,
"docker_rootless": self.docker_rootless,
"nvidia_container_runtime": self.nvidia_container_runtime,
"gpus": [gpu.to_dict() for gpu in self.gpus],
"probed_at_epoch": self.probed_at_epoch,
"reported_cost_usd_per_second": self.reported_cost_usd_per_second,
}
@classmethod
def from_dict(cls, payload: Mapping[str, Any]) -> "EnvironmentCapabilities":
"""Build validated capabilities from persisted data."""
raw_gpus = payload.get("gpus")
gpu_payloads = raw_gpus if isinstance(raw_gpus, list) else []
return cls(
architecture=str(payload.get("architecture") or "unknown").strip(),
operating_system=str(payload.get("operating_system") or "unknown").strip(),
cpu_count=_non_negative_int(payload.get("cpu_count"), "cpu_count"),
total_ram_bytes=_non_negative_int(
payload.get("total_ram_bytes"), "total_ram_bytes"
),
available_ram_bytes=_optional_non_negative_int(
payload.get("available_ram_bytes"), "available_ram_bytes"
),
available_disk_bytes=_optional_non_negative_int(
payload.get("available_disk_bytes"), "available_disk_bytes"
),
docker_version=str(payload.get("docker_version") or "unknown").strip(),
docker_rootless=bool(payload.get("docker_rootless", False)),
nvidia_container_runtime=bool(
payload.get("nvidia_container_runtime", False)
),
gpus=tuple(
GpuCapability.from_dict(gpu_payload)
for gpu_payload in gpu_payloads
if isinstance(gpu_payload, Mapping)
),
probed_at_epoch=_optional_non_negative_float(
payload.get("probed_at_epoch"), "probed_at_epoch"
),
reported_cost_usd_per_second=_optional_non_negative_float(
payload.get("reported_cost_usd_per_second"),
"reported_cost_usd_per_second",
),
)
@dataclass(frozen=True)
class ComponentResourceRequirements:
"""Describe hard requirements and cost inputs for one remote component."""
minimum_vram_bytes: int = 0
minimum_ram_bytes: int = 0
gpu_required: bool = True
architecture: str | None = None
estimated_execution_seconds: float = 0.0
estimated_execution_seconds_by_environment: Mapping[str, float] = field(
default_factory=dict
)
estimated_transfer_seconds: float = 0.0
required_tags: frozenset[str] = frozenset()
preferred_environment_ids: tuple[str, ...] = ()
required_provider: ExecutionProvider | None = None
@dataclass(frozen=True)
class EnvironmentSchedulingState:
"""Combine configured and transient state used for scheduling."""
environment_id: str
provider: ExecutionProvider
enabled: bool
health: EnvironmentHealth
cost_usd_per_second: float | None
capabilities: EnvironmentCapabilities | None
configuration_id: str | None = None
display_name: str | None = None
unavailable_reason: str | None = None
tags: frozenset[str] = frozenset()
active_workers: int = 0
maximum_workers: int = 1
queue_delay_seconds: float = 0.0
cold_start_seconds: float = 0.0
cached_artifact_keys: frozenset[str] = frozenset()
@dataclass(frozen=True)
class ExecutionAssignment:
"""Record the selected environment and its predicted scheduling cost."""
environment_id: str
provider: ExecutionProvider
predicted_cost_usd: float | None
predicted_completion_seconds: float
reasons: tuple[str, ...] = ()
configuration_id: str | None = None
capacity_slot_index: int = 0
@dataclass(frozen=True)
class WorkflowExecutionPreferences:
"""Describe provider and placement preferences saved with one workflow."""
policy: ExecutionPolicy = ExecutionPolicy.MODAL
auto_place: bool = False
preferred_environment_ids: tuple[str, ...] = ()
minimum_vram_bytes: int = 0
@classmethod
def from_workflow(cls, workflow: object) -> "WorkflowExecutionPreferences":
"""Read generic preferences while preserving legacy Modal defaults."""
if not isinstance(workflow, Mapping):
return cls()
extra = workflow.get("extra")
if not isinstance(extra, Mapping):
return cls()
payload = extra.get(WORKFLOW_REMOTE_CONFIG_KEY)
if not isinstance(payload, Mapping):
return cls()
raw_preferred = payload.get(WORKFLOW_REMOTE_PREFERRED_ENVIRONMENTS_KEY)
preferred_values = raw_preferred if isinstance(raw_preferred, list) else []
raw_minimum_vram_gb = payload.get(WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY, 0)
minimum_vram_gb = _optional_non_negative_float(
raw_minimum_vram_gb,
WORKFLOW_REMOTE_MINIMUM_VRAM_GB_KEY,
)
return cls(
policy=ExecutionPolicy(
str(payload.get(WORKFLOW_REMOTE_POLICY_KEY) or "modal").strip().lower()
),
auto_place=bool(payload.get(WORKFLOW_REMOTE_AUTO_PLACE_KEY, False)),
preferred_environment_ids=tuple(
str(value).strip() for value in preferred_values if str(value).strip()
),
minimum_vram_bytes=int((minimum_vram_gb or 0.0) * 1024**3),
)
class ExecutionEnvironmentClient(Protocol):
"""Provider-neutral client used to prepare and invoke remote environments."""
def execute_payload(
self,
payload: Mapping[str, Any],
kwargs: Mapping[str, Any],
) -> Sequence[Any]:
"""Execute one component payload and return deserialized outputs."""
async def execute_payload_async(
self,
payload: Mapping[str, Any],
kwargs: Mapping[str, Any],
) -> Sequence[Any]:
"""Execute one component payload asynchronously."""
class NoCompatibleExecutionEnvironmentError(RuntimeError):
"""Raised when no configured environment satisfies component requirements."""
class CostAwareEnvironmentScheduler:
"""Select the least-cost compatible environment with deterministic tie-breaking."""
def choose(
self,
environments: Sequence[EnvironmentSchedulingState],
requirements: ComponentResourceRequirements,
) -> ExecutionAssignment:
"""Return the best currently compatible execution environment."""
candidates: list[tuple[tuple[Any, ...], ExecutionAssignment]] = []
rejection_reasons: list[str] = []
for environment in environments:
rejection = self._incompatibility_reason(environment, requirements)
if rejection is not None:
rejection_reasons.append(f"{environment.environment_id}: {rejection}")
continue
execution_seconds = (
requirements.estimated_execution_seconds_by_environment.get(
environment.environment_id,
requirements.estimated_execution_seconds,
)
)
completion_seconds = max(
0.0,
environment.queue_delay_seconds
+ environment.cold_start_seconds
+ requirements.estimated_transfer_seconds
+ execution_seconds,
)
predicted_cost = (
completion_seconds * environment.cost_usd_per_second
if environment.cost_usd_per_second is not None
else None
)
preferred_rank = self._preferred_rank(
environment.environment_id,
requirements.preferred_environment_ids,
)
unknown_cost_rank = 1 if predicted_cost is None else 0
cost_rank = predicted_cost if predicted_cost is not None else math.inf
assignment = ExecutionAssignment(
environment_id=environment.environment_id,
provider=environment.provider,
predicted_cost_usd=predicted_cost,
predicted_completion_seconds=completion_seconds,
reasons=(
f"compatible with {len(environment.capabilities.gpus) if environment.capabilities else 0} GPU(s)",
(
"requires at least "
f"{requirements.minimum_vram_bytes / 1024**3:.2f} GiB "
"GPU VRAM"
if requirements.minimum_vram_bytes
else "no explicit GPU VRAM floor"
),
(
f"estimated cost ${predicted_cost:.6f}"
if predicted_cost is not None
else "cost unknown"
),
f"estimated execution {execution_seconds:.3f}s",
),
)
candidates.append(
(
(
preferred_rank,
unknown_cost_rank,
cost_rank,
completion_seconds,
environment.environment_id,
),
assignment,
)
)
if not candidates:
details = "; ".join(rejection_reasons) or "no environments are configured"
raise NoCompatibleExecutionEnvironmentError(
f"No compatible remote execution environment is available: {details}."
)
return min(candidates, key=lambda candidate: candidate[0])[1]
def plan(
self,
*,
execution_stages: Sequence[Sequence[str]],
environments_by_component: Mapping[
str,
Sequence[EnvironmentSchedulingState],
],
requirements_by_component: Mapping[str, ComponentResourceRequirements],
) -> dict[str, ExecutionAssignment]:
"""Assign a staged component DAG while honoring pool capacity limits."""
slot_available_at: dict[tuple[str, int], float] = {}
assignments: dict[str, ExecutionAssignment] = {}
stage_ready_at = 0.0
for stage in execution_stages:
stage_completion_times: list[float] = []
for component_id in sorted(str(value) for value in stage):
requirements = requirements_by_component.get(component_id)
if requirements is None:
raise ValueError(
f"Remote component {component_id!r} has no resource "
"requirements."
)
environments = environments_by_component.get(component_id, ())
candidates: list[tuple[tuple[Any, ...], ExecutionAssignment]] = []
rejection_reasons: list[str] = []
for environment in environments:
rejection = self._incompatibility_reason(
environment,
requirements,
)
if rejection is not None:
rejection_reasons.append(
f"{environment.display_name or environment.environment_id}: "
f"{rejection}"
)
continue
for slot_index in range(environment.maximum_workers):
slot_key = (environment.environment_id, slot_index)
slot_delay = max(
0.0,
slot_available_at.get(slot_key, stage_ready_at)
- stage_ready_at,
)
slot_environment = EnvironmentSchedulingState(
environment_id=environment.environment_id,
provider=environment.provider,
enabled=environment.enabled,
health=environment.health,
cost_usd_per_second=environment.cost_usd_per_second,
capabilities=environment.capabilities,
configuration_id=environment.configuration_id,
display_name=environment.display_name,
unavailable_reason=environment.unavailable_reason,
tags=environment.tags,
active_workers=0,
maximum_workers=1,
queue_delay_seconds=(
environment.queue_delay_seconds + slot_delay
),
cold_start_seconds=(
environment.cold_start_seconds
if slot_key not in slot_available_at
else 0.0
),
cached_artifact_keys=environment.cached_artifact_keys,
)
assignment = self.choose([slot_environment], requirements)
absolute_completion = (
stage_ready_at
+ assignment.predicted_completion_seconds
)
configuration_id = (
environment.configuration_id
or environment.environment_id
)
assignment = ExecutionAssignment(
environment_id=assignment.environment_id,
provider=assignment.provider,
predicted_cost_usd=assignment.predicted_cost_usd,
predicted_completion_seconds=absolute_completion,
reasons=assignment.reasons
+ (
"configuration "
f"{environment.display_name or configuration_id}",
f"capacity slot {slot_index}",
),
configuration_id=configuration_id,
capacity_slot_index=slot_index,
)
preferred_rank = self._preferred_rank(
environment.environment_id,
requirements.preferred_environment_ids,
)
unknown_cost_rank = (
1 if assignment.predicted_cost_usd is None else 0
)
cost_rank = (
assignment.predicted_cost_usd
if assignment.predicted_cost_usd is not None
else math.inf
)
candidates.append(
(
(
preferred_rank,
unknown_cost_rank,
cost_rank,
absolute_completion,
environment.environment_id,
slot_index,
),
assignment,
)
)
if not candidates:
details = (
"; ".join(rejection_reasons)
or "no environments are configured"
)
raise NoCompatibleExecutionEnvironmentError(
"No compatible remote execution environment is available "
f"for component {component_id!r}: {details}."
)
assignment = min(candidates, key=lambda candidate: candidate[0])[1]
assignments[component_id] = assignment
slot_available_at[
(assignment.environment_id, assignment.capacity_slot_index)
] = assignment.predicted_completion_seconds
stage_completion_times.append(
assignment.predicted_completion_seconds
)
if stage_completion_times:
stage_ready_at = max(stage_completion_times)
return assignments
def _incompatibility_reason(
self,
environment: EnvironmentSchedulingState,
requirements: ComponentResourceRequirements,
) -> str | None:
"""Return a human-readable hard-constraint failure when incompatible."""
if not environment.enabled:
return "disabled"
if (
requirements.required_provider is not None
and environment.provider is not requirements.required_provider
):
return f"provider is not {requirements.required_provider.value!r}"
if environment.health not in {
EnvironmentHealth.READY,
EnvironmentHealth.DEGRADED,
}:
if environment.unavailable_reason:
return environment.unavailable_reason
return f"health is {environment.health.value}"
if environment.maximum_workers <= 0:
return "worker limit is zero"
if environment.active_workers >= environment.maximum_workers:
return "all worker capacity is active"
capabilities = environment.capabilities
if capabilities is None:
return "capabilities have not been probed"
if (
requirements.architecture
and capabilities.architecture != requirements.architecture
):
return (
f"architecture {capabilities.architecture!r} does not satisfy "
f"{requirements.architecture!r}"
)
if capabilities.schedulable_ram_bytes < requirements.minimum_ram_bytes:
return (
"insufficient RAM "
f"({capabilities.schedulable_ram_bytes / 1024**3:.2f} GiB available, "
f"{requirements.minimum_ram_bytes / 1024**3:.2f} GiB required)"
)
if not requirements.required_tags.issubset(environment.tags):
return "required tags are missing"
if requirements.gpu_required:
if not capabilities.nvidia_container_runtime:
return "NVIDIA container runtime is unavailable"
if (
capabilities.maximum_available_vram_bytes
< requirements.minimum_vram_bytes
):
return (
"insufficient GPU VRAM "
"("
f"{capabilities.maximum_available_vram_bytes / 1024**3:.2f} "
"GiB available, "
f"{requirements.minimum_vram_bytes / 1024**3:.2f} GiB required)"
)
return None
def _preferred_rank(
self,
environment_id: str,
preferred_environment_ids: tuple[str, ...],
) -> int:
"""Return a stable rank for explicit environment preferences."""
if not preferred_environment_ids:
return 0
try:
return preferred_environment_ids.index(environment_id)
except ValueError:
return len(preferred_environment_ids) + 1
def _non_negative_int(value: Any, field_name: str) -> int:
"""Return one non-negative integer or raise a validation error."""
if isinstance(value, bool):
raise ValueError(f"{field_name} must be a non-negative integer.")
try:
normalized = int(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"{field_name} must be a non-negative integer.") from exc
if normalized < 0:
raise ValueError(f"{field_name} must be a non-negative integer.")
return normalized
def _optional_non_negative_int(value: Any, field_name: str) -> int | None:
"""Return an optional non-negative integer."""
if value is None:
return None
return _non_negative_int(value, field_name)
def _optional_non_negative_float(value: Any, field_name: str) -> float | None:
"""Return an optional finite non-negative floating-point number."""
if value is None:
return None
if isinstance(value, bool):
raise ValueError(f"{field_name} must be a non-negative number.")
try:
normalized = float(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"{field_name} must be a non-negative number.") from exc
if not math.isfinite(normalized) or normalized < 0:
raise ValueError(f"{field_name} must be a non-negative number.")
return normalized
def _optional_string(value: Any) -> str | None:
"""Return stripped optional text."""
if value is None:
return None
normalized = str(value).strip()
return normalized or None