1433 lines
54 KiB
Python
1433 lines
54 KiB
Python
"""Provider-neutral execution planning, capacity acquisition, and assignment metadata."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import defaultdict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping
|
|
|
|
if __package__:
|
|
from .execution_assignment_runtime import (
|
|
_configured_provider_metadata,
|
|
_ensure_remote_sync_backend,
|
|
_execution_location_for_assignment,
|
|
_ssh_hostname,
|
|
_ssh_sync_engine,
|
|
_stamp_execution_assignment,
|
|
_vast_provider_metadata,
|
|
_workflow_r2_cache,
|
|
)
|
|
from .execution_environments import (
|
|
ComponentResourceRequirements,
|
|
CostAwareEnvironmentScheduler,
|
|
EnvironmentCapabilities,
|
|
EnvironmentHealth,
|
|
EnvironmentSchedulingState,
|
|
ExecutionAssignment,
|
|
ExecutionPolicy,
|
|
ExecutionProvider,
|
|
GpuCapability,
|
|
NoCompatibleExecutionEnvironmentError,
|
|
WorkflowExecutionPreferences,
|
|
)
|
|
from .execution_history import ExecutionHistory
|
|
from .execution_plan_reporting import (
|
|
_assignment_hardware_payload,
|
|
_cached_r2_storage_usage,
|
|
_configuration_field,
|
|
_configuration_host,
|
|
_planned_execution_assignments_payload,
|
|
_r2_storage_from_usage_payload,
|
|
_refresh_r2_storage_usage,
|
|
_safe_remote_configuration_payload,
|
|
)
|
|
from .execution_resource_estimates import (
|
|
_component_execution_signature,
|
|
_component_memory_estimate,
|
|
_component_model_asset_sizes,
|
|
_component_profile_memory_estimate,
|
|
_component_required_provider,
|
|
_is_additive_model_node,
|
|
_iter_prompt_string_values,
|
|
_prompt_llm_model_references,
|
|
_resolve_prompt_llm_profiles,
|
|
)
|
|
from .modal_hardware import (
|
|
_MODAL_GPU_COST_USD_PER_SECOND,
|
|
_MODAL_GPU_VRAM_GB,
|
|
)
|
|
from .remote_configuration_nodes import compile_remote_configuration_set
|
|
from .remote_configurations import (
|
|
ModalRemoteConfiguration,
|
|
RemoteConfiguration,
|
|
RemoteConfigurationSet,
|
|
SshRemoteConfiguration,
|
|
VastRemoteConfiguration,
|
|
)
|
|
from .remote_graph_analysis import _component_execution_stages
|
|
from .remote_hosts import RemoteHostRegistry, SshHostConfig
|
|
from .remote_plan_types import (
|
|
ComponentExecutionPlan,
|
|
ComponentMemoryEstimate,
|
|
ModalPromptValidationError,
|
|
RemoteComponentPlan,
|
|
)
|
|
from .settings import (
|
|
ModalSyncSettings,
|
|
discover_comfyui_user_directory,
|
|
settings_for_modal_gpu,
|
|
)
|
|
from .ssh_docker import SshDockerController
|
|
from .sync_engine import SyncCancelledError
|
|
from .vast_config_node import extract_vast_profiles
|
|
from .vast_service import VastProfileQuote, VastSearchRequirements, VastService
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from execution_assignment_runtime import (
|
|
_configured_provider_metadata,
|
|
_ensure_remote_sync_backend,
|
|
_execution_location_for_assignment,
|
|
_ssh_hostname,
|
|
_ssh_sync_engine,
|
|
_stamp_execution_assignment,
|
|
_vast_provider_metadata,
|
|
_workflow_r2_cache,
|
|
)
|
|
from execution_environments import (
|
|
ComponentResourceRequirements,
|
|
CostAwareEnvironmentScheduler,
|
|
EnvironmentCapabilities,
|
|
EnvironmentHealth,
|
|
EnvironmentSchedulingState,
|
|
ExecutionAssignment,
|
|
ExecutionPolicy,
|
|
ExecutionProvider,
|
|
GpuCapability,
|
|
NoCompatibleExecutionEnvironmentError,
|
|
WorkflowExecutionPreferences,
|
|
)
|
|
from execution_history import ExecutionHistory
|
|
from execution_plan_reporting import (
|
|
_assignment_hardware_payload,
|
|
_cached_r2_storage_usage,
|
|
_configuration_field,
|
|
_configuration_host,
|
|
_planned_execution_assignments_payload,
|
|
_r2_storage_from_usage_payload,
|
|
_refresh_r2_storage_usage,
|
|
_safe_remote_configuration_payload,
|
|
)
|
|
from execution_resource_estimates import (
|
|
_component_execution_signature,
|
|
_component_memory_estimate,
|
|
_component_model_asset_sizes,
|
|
_component_profile_memory_estimate,
|
|
_component_required_provider,
|
|
_is_additive_model_node,
|
|
_iter_prompt_string_values,
|
|
_prompt_llm_model_references,
|
|
_resolve_prompt_llm_profiles,
|
|
)
|
|
from modal_hardware import (
|
|
_MODAL_GPU_COST_USD_PER_SECOND,
|
|
_MODAL_GPU_VRAM_GB,
|
|
)
|
|
from remote_configuration_nodes import compile_remote_configuration_set
|
|
from remote_configurations import (
|
|
ModalRemoteConfiguration,
|
|
RemoteConfiguration,
|
|
RemoteConfigurationSet,
|
|
SshRemoteConfiguration,
|
|
VastRemoteConfiguration,
|
|
)
|
|
from remote_graph_analysis import _component_execution_stages
|
|
from remote_hosts import RemoteHostRegistry, SshHostConfig
|
|
from remote_plan_types import (
|
|
ComponentExecutionPlan,
|
|
ComponentMemoryEstimate,
|
|
ModalPromptValidationError,
|
|
RemoteComponentPlan,
|
|
)
|
|
from settings import (
|
|
ModalSyncSettings,
|
|
discover_comfyui_user_directory,
|
|
settings_for_modal_gpu,
|
|
)
|
|
from ssh_docker import SshDockerController
|
|
from sync_engine import SyncCancelledError
|
|
from vast_config_node import extract_vast_profiles
|
|
from vast_service import VastProfileQuote, VastSearchRequirements, VastService
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SetupStatusCallback = Callable[[str, int | None, int | None], None]
|
|
EnvironmentSetupStatusCallback = Callable[
|
|
[str, str, int | None, int | None], None
|
|
]
|
|
ExecutionPlanStatusCallback = Callable[
|
|
[dict[str, dict[str, Any]], list[dict[str, Any]]], None
|
|
]
|
|
|
|
def _modal_environment_state(settings: ModalSyncSettings) -> EnvironmentSchedulingState:
|
|
"""Return the scheduler-facing state of the selected Modal GPU target."""
|
|
modal_gpu = settings.modal_gpu
|
|
vram_bytes = int(_MODAL_GPU_VRAM_GB.get(modal_gpu, 0.0) * 1024**3)
|
|
capabilities = EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=1,
|
|
total_ram_bytes=max(vram_bytes, 1),
|
|
available_ram_bytes=None,
|
|
available_disk_bytes=None,
|
|
docker_version="modal-managed",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=(
|
|
GpuCapability(
|
|
uuid=f"modal-{modal_gpu.lower()}",
|
|
name=modal_gpu,
|
|
total_vram_bytes=vram_bytes,
|
|
),
|
|
),
|
|
)
|
|
return EnvironmentSchedulingState(
|
|
environment_id=f"modal:{modal_gpu}",
|
|
provider=ExecutionProvider.MODAL,
|
|
enabled=True,
|
|
health=EnvironmentHealth.READY,
|
|
cost_usd_per_second=_MODAL_GPU_COST_USD_PER_SECOND.get(modal_gpu),
|
|
capabilities=capabilities,
|
|
maximum_workers=settings.max_containers or 1,
|
|
)
|
|
|
|
|
|
def _modal_configuration_environment_state(
|
|
configuration: ModalRemoteConfiguration,
|
|
settings: ModalSyncSettings,
|
|
) -> EnvironmentSchedulingState:
|
|
"""Return one scheduler state for a workflow-declared Modal capacity pool."""
|
|
selected_settings = settings_for_modal_gpu(settings, configuration.gpu_type)
|
|
base_state = _modal_environment_state(selected_settings)
|
|
return replace(
|
|
base_state,
|
|
environment_id=(
|
|
f"modal:{configuration.configuration_id}:{configuration.gpu_type}"
|
|
),
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
maximum_workers=configuration.instance_count,
|
|
)
|
|
|
|
|
|
def _ssh_host_registry(settings: ModalSyncSettings) -> RemoteHostRegistry | None:
|
|
"""Return the persistent SSH host registry when a user directory exists."""
|
|
user_directory = discover_comfyui_user_directory(settings)
|
|
if user_directory is None:
|
|
return None
|
|
return RemoteHostRegistry.for_user_directory(user_directory)
|
|
|
|
|
|
def _configured_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, ...]:
|
|
"""Return configured SSH hosts without probing during queue submission."""
|
|
registry = _ssh_host_registry(settings)
|
|
if registry is None:
|
|
return ()
|
|
return registry.load().hosts
|
|
|
|
|
|
def _schedulable_ssh_hosts(settings: ModalSyncSettings) -> tuple[SshHostConfig, ...]:
|
|
"""Probe every enabled host immediately before cost-aware scheduling."""
|
|
hosts = _configured_ssh_hosts(settings)
|
|
registry = _ssh_host_registry(settings)
|
|
if registry is None:
|
|
return hosts
|
|
probe_hosts = [host for host in hosts if host.enabled and not host.draining]
|
|
if not probe_hosts:
|
|
return hosts
|
|
|
|
with ThreadPoolExecutor(max_workers=min(8, len(probe_hosts))) as executor:
|
|
refreshed = {
|
|
host.environment_id: future.result()
|
|
for host, future in (
|
|
(host, executor.submit(_refresh_ssh_host, host, registry))
|
|
for host in probe_hosts
|
|
)
|
|
}
|
|
return tuple(refreshed.get(host.environment_id, host) for host in hosts)
|
|
|
|
|
|
def _refresh_ssh_host(
|
|
host: SshHostConfig,
|
|
registry: RemoteHostRegistry,
|
|
) -> SshHostConfig:
|
|
"""Probe and persist one host without aborting other candidates."""
|
|
try:
|
|
capabilities = SshDockerController(host).probe_capabilities()
|
|
return registry.update_probe_result(
|
|
host.environment_id,
|
|
capabilities=capabilities,
|
|
health=EnvironmentHealth.READY,
|
|
last_error=None,
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Remote environment probe failed environment=%s: %s",
|
|
host.environment_id,
|
|
exc,
|
|
)
|
|
try:
|
|
return registry.update_probe_result(
|
|
host.environment_id,
|
|
capabilities=host.capabilities,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
last_error=str(exc),
|
|
)
|
|
except (KeyError, OSError, ValueError):
|
|
return replace(
|
|
host,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
last_error=str(exc),
|
|
)
|
|
|
|
|
|
def _execution_history(settings: ModalSyncSettings) -> ExecutionHistory | None:
|
|
"""Return local runtime history when the ComfyUI user directory is known."""
|
|
user_directory = discover_comfyui_user_directory(settings)
|
|
if user_directory is None:
|
|
return None
|
|
return ExecutionHistory.for_user_directory(user_directory)
|
|
|
|
|
|
def _maximum_capacity_state(
|
|
environment: EnvironmentSchedulingState,
|
|
) -> EnvironmentSchedulingState:
|
|
"""Return an optimistic state used only to identify reclaimable SSH capacity."""
|
|
capabilities = environment.capabilities
|
|
if environment.provider is not ExecutionProvider.SSH_DOCKER or capabilities is None:
|
|
return environment
|
|
return replace(
|
|
environment,
|
|
capabilities=replace(
|
|
capabilities,
|
|
available_ram_bytes=capabilities.total_ram_bytes,
|
|
gpus=tuple(
|
|
replace(gpu, free_vram_bytes=gpu.total_vram_bytes)
|
|
for gpu in capabilities.gpus
|
|
),
|
|
),
|
|
)
|
|
|
|
|
|
def _reprobe_reclaimed_ssh_host(
|
|
host: SshHostConfig,
|
|
settings: ModalSyncSettings,
|
|
) -> SshHostConfig:
|
|
"""Probe a reclaimed host and persist its new free-memory measurements."""
|
|
registry = _ssh_host_registry(settings)
|
|
if registry is not None:
|
|
return _refresh_ssh_host(host, registry)
|
|
try:
|
|
capabilities = SshDockerController(host).probe_capabilities()
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Remote environment re-probe failed after worker reclaim "
|
|
"environment=%s: %s",
|
|
host.environment_id,
|
|
exc,
|
|
)
|
|
return replace(
|
|
host,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
last_error=str(exc),
|
|
)
|
|
return replace(
|
|
host,
|
|
capabilities=capabilities,
|
|
health=EnvironmentHealth.READY,
|
|
last_error=None,
|
|
)
|
|
|
|
|
|
def _probe_workflow_ssh_configuration(
|
|
configuration: SshRemoteConfiguration,
|
|
) -> SshHostConfig:
|
|
"""Probe one workflow-declared SSH host without mutating the global registry."""
|
|
host = configuration.host
|
|
try:
|
|
capabilities = SshDockerController(host).probe_capabilities()
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Workflow SSH environment probe failed configuration=%s environment=%s: %s",
|
|
configuration.display_name,
|
|
host.environment_id,
|
|
exc,
|
|
)
|
|
return replace(
|
|
host,
|
|
capabilities=None,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
last_error=str(exc),
|
|
)
|
|
return replace(
|
|
host,
|
|
capabilities=capabilities,
|
|
health=EnvironmentHealth.READY,
|
|
last_error=None,
|
|
)
|
|
|
|
|
|
def _optional_scheduler_choice(
|
|
scheduler: CostAwareEnvironmentScheduler,
|
|
environments: list[EnvironmentSchedulingState],
|
|
requirements: ComponentResourceRequirements,
|
|
) -> tuple[
|
|
ExecutionAssignment | None,
|
|
NoCompatibleExecutionEnvironmentError | None,
|
|
]:
|
|
"""Return a scheduler choice and preserve a compatibility failure for fallback."""
|
|
try:
|
|
return scheduler.choose(environments, requirements), None
|
|
except NoCompatibleExecutionEnvironmentError as exc:
|
|
return None, exc
|
|
|
|
|
|
def _require_scheduler_choice(
|
|
assignment: ExecutionAssignment | None,
|
|
error: NoCompatibleExecutionEnvironmentError | None,
|
|
) -> ExecutionAssignment:
|
|
"""Return an existing scheduler choice or raise its compatibility detail."""
|
|
if assignment is not None:
|
|
return assignment
|
|
if error is not None:
|
|
raise error
|
|
raise RuntimeError("Scheduler failed without an incompatibility reason.")
|
|
|
|
|
|
def _remove_idle_ssh_workers_for_reclaim(host: SshHostConfig) -> tuple[str, ...]:
|
|
"""Recycle idle managed workers without disrupting fallback placement."""
|
|
try:
|
|
return SshDockerController(host).remove_idle_managed_workers()
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Could not reclaim managed SSH worker capacity environment=%s: %s",
|
|
host.environment_id,
|
|
exc,
|
|
)
|
|
return ()
|
|
|
|
|
|
def _reclaim_improves_assignment(
|
|
optimistic: ExecutionAssignment,
|
|
actual: ExecutionAssignment | None,
|
|
requirements: ComponentResourceRequirements,
|
|
) -> bool:
|
|
"""Return whether reclaim unlocks required or materially better placement."""
|
|
if actual is None:
|
|
return True
|
|
if optimistic.environment_id == actual.environment_id:
|
|
return False
|
|
preferred = requirements.preferred_environment_ids
|
|
if optimistic.environment_id in preferred:
|
|
if actual.environment_id not in preferred:
|
|
return True
|
|
return preferred.index(optimistic.environment_id) < preferred.index(
|
|
actual.environment_id
|
|
)
|
|
if optimistic.predicted_cost_usd is None:
|
|
return False
|
|
if actual.predicted_cost_usd is None:
|
|
return True
|
|
return optimistic.predicted_cost_usd < actual.predicted_cost_usd
|
|
|
|
|
|
def _choose_with_idle_ssh_worker_reclaim(
|
|
*,
|
|
scheduler: CostAwareEnvironmentScheduler,
|
|
environments: list[EnvironmentSchedulingState],
|
|
requirements: ComponentResourceRequirements,
|
|
ssh_hosts_by_id: dict[str, SshHostConfig],
|
|
settings: ModalSyncSettings,
|
|
) -> ExecutionAssignment:
|
|
"""Choose an environment, recycling a cheaper idle SSH worker when necessary."""
|
|
actual_assignment, actual_error = _optional_scheduler_choice(
|
|
scheduler,
|
|
environments,
|
|
requirements,
|
|
)
|
|
optimistic_assignment, _ = _optional_scheduler_choice(
|
|
scheduler,
|
|
[_maximum_capacity_state(environment) for environment in environments],
|
|
requirements,
|
|
)
|
|
if optimistic_assignment is None:
|
|
return _require_scheduler_choice(actual_assignment, actual_error)
|
|
|
|
if (
|
|
optimistic_assignment.provider is not ExecutionProvider.SSH_DOCKER
|
|
or not _reclaim_improves_assignment(
|
|
optimistic_assignment,
|
|
actual_assignment,
|
|
requirements,
|
|
)
|
|
):
|
|
return _require_scheduler_choice(actual_assignment, actual_error)
|
|
|
|
host = ssh_hosts_by_id.get(optimistic_assignment.environment_id)
|
|
if host is None:
|
|
return _require_scheduler_choice(actual_assignment, actual_error)
|
|
|
|
removed_workers = _remove_idle_ssh_workers_for_reclaim(host)
|
|
if not removed_workers:
|
|
return _require_scheduler_choice(actual_assignment, actual_error)
|
|
|
|
logger.info(
|
|
"Recycled idle managed SSH worker(s) environment=%s containers=%s so the "
|
|
"planner can re-evaluate full host capacity.",
|
|
host.environment_id,
|
|
removed_workers,
|
|
)
|
|
refreshed_host = _reprobe_reclaimed_ssh_host(host, settings)
|
|
ssh_hosts_by_id[host.environment_id] = refreshed_host
|
|
refreshed_environments = [
|
|
(
|
|
refreshed_host.scheduling_state()
|
|
if environment.environment_id == host.environment_id
|
|
else environment
|
|
)
|
|
for environment in environments
|
|
]
|
|
return scheduler.choose(refreshed_environments, requirements)
|
|
|
|
|
|
def _plan_component_execution_assignments(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
prompt: Mapping[str, Any],
|
|
workflow: Mapping[str, Any] | None,
|
|
settings: ModalSyncSettings,
|
|
status_callback: SetupStatusCallback | None = None,
|
|
resolved_llm_profiles: Mapping[str, Any] | None = None,
|
|
) -> dict[str, ExecutionAssignment]:
|
|
"""Assign components across Modal, SSH Docker, and workflow-declared Vast pools."""
|
|
preferences = WorkflowExecutionPreferences.from_workflow(workflow)
|
|
active_llm_profiles = (
|
|
dict(resolved_llm_profiles)
|
|
if resolved_llm_profiles is not None
|
|
else _resolve_prompt_llm_profiles(prompt, settings)
|
|
)
|
|
if preferences.policy is ExecutionPolicy.MODAL:
|
|
incompatible_component_ids = [
|
|
component.representative_node_id
|
|
for component in components
|
|
if _component_required_provider(
|
|
component,
|
|
prompt,
|
|
active_llm_profiles,
|
|
)
|
|
is ExecutionProvider.SSH_DOCKER
|
|
]
|
|
if incompatible_component_ids:
|
|
raise ModalPromptValidationError(
|
|
"Modal-only execution cannot run SSH-only component(s) "
|
|
f"{incompatible_component_ids}. Select Automatic or Self-hosted "
|
|
"execution for the workflow."
|
|
)
|
|
modal_state = _modal_environment_state(settings)
|
|
return {
|
|
component.representative_node_id: ExecutionAssignment(
|
|
environment_id=modal_state.environment_id,
|
|
provider=ExecutionProvider.MODAL,
|
|
predicted_cost_usd=None,
|
|
predicted_completion_seconds=0.0,
|
|
reasons=("workflow policy requires Modal",),
|
|
)
|
|
for component in components
|
|
}
|
|
|
|
try:
|
|
vast_profiles = extract_vast_profiles(prompt)
|
|
except ValueError as exc:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
vast_service: VastService | None = None
|
|
if preferences.policy in {ExecutionPolicy.VAST, ExecutionPolicy.AUTOMATIC}:
|
|
if preferences.policy is ExecutionPolicy.VAST and not vast_profiles:
|
|
raise ModalPromptValidationError(
|
|
"Vast.ai-only execution requires at least one disconnected "
|
|
"Vast.ai Lease Configuration node in the workflow."
|
|
)
|
|
if vast_profiles:
|
|
try:
|
|
vast_service = VastService.from_environment(
|
|
settings,
|
|
repo_root=Path(__file__).resolve().parent,
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
if preferences.policy is ExecutionPolicy.VAST:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
logger.warning("Skipping Vast.ai automatic placement: %s", exc)
|
|
|
|
ssh_hosts = (
|
|
_schedulable_ssh_hosts(settings)
|
|
if preferences.policy in {ExecutionPolicy.SELF_HOSTED, ExecutionPolicy.AUTOMATIC}
|
|
else ()
|
|
)
|
|
ssh_hosts_by_id = {host.environment_id: host for host in ssh_hosts}
|
|
modal_state = (
|
|
_modal_environment_state(settings)
|
|
if preferences.policy is ExecutionPolicy.AUTOMATIC
|
|
else None
|
|
)
|
|
if not ssh_hosts and modal_state is None and vast_service is None:
|
|
raise ModalPromptValidationError(
|
|
"This workflow requests remote execution, but no compatible provider is configured."
|
|
)
|
|
|
|
scheduler = CostAwareEnvironmentScheduler()
|
|
history = _execution_history(settings)
|
|
component_signatures: dict[str, str] = {}
|
|
component_memory_estimates: dict[str, ComponentMemoryEstimate] = {}
|
|
component_required_providers: dict[str, ExecutionProvider | None] = {}
|
|
for component in components:
|
|
component_id = component.representative_node_id
|
|
component_signatures[component_id] = _component_execution_signature(
|
|
component,
|
|
prompt,
|
|
)
|
|
memory_estimate = _component_memory_estimate(
|
|
component,
|
|
prompt,
|
|
preferences,
|
|
settings,
|
|
active_llm_profiles,
|
|
)
|
|
component_memory_estimates[component_id] = memory_estimate
|
|
component_required_providers[component_id] = _component_required_provider(
|
|
component,
|
|
prompt,
|
|
active_llm_profiles,
|
|
)
|
|
if memory_estimate.model_asset_count:
|
|
logger.info(
|
|
"Estimated remote component=%s memory floor vram_gib=%.2f ram_gib=%.2f "
|
|
"from model_assets=%d largest_model_gib=%.2f.",
|
|
component_id,
|
|
memory_estimate.minimum_vram_bytes / 1024**3,
|
|
memory_estimate.minimum_ram_bytes / 1024**3,
|
|
memory_estimate.model_asset_count,
|
|
memory_estimate.largest_model_bytes / 1024**3,
|
|
)
|
|
if vast_service is not None:
|
|
try:
|
|
vast_service.prefetch_offers_sync(
|
|
vast_profiles,
|
|
tuple(
|
|
VastSearchRequirements(
|
|
minimum_vram_bytes=memory_estimate.minimum_vram_bytes,
|
|
minimum_ram_bytes=memory_estimate.minimum_ram_bytes,
|
|
)
|
|
for memory_estimate in component_memory_estimates.values()
|
|
),
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
if preferences.policy is ExecutionPolicy.VAST:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
logger.warning(
|
|
"Unable to prefetch Vast.ai marketplace offers: %s",
|
|
exc,
|
|
)
|
|
assignments: dict[str, ExecutionAssignment] = {}
|
|
for component in components:
|
|
component_id = component.representative_node_id
|
|
environments = [
|
|
host.scheduling_state() for host in ssh_hosts_by_id.values()
|
|
]
|
|
if modal_state is not None:
|
|
environments.append(modal_state)
|
|
component_signature = component_signatures[component_id]
|
|
memory_estimate = component_memory_estimates[component_id]
|
|
vast_quote: VastProfileQuote | None = None
|
|
if vast_service is not None:
|
|
try:
|
|
vast_quote = vast_service.quote_best_profile_sync(
|
|
vast_profiles,
|
|
minimum_vram_bytes=memory_estimate.minimum_vram_bytes,
|
|
minimum_ram_bytes=memory_estimate.minimum_ram_bytes,
|
|
predicted_execution_seconds=60.0,
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
if preferences.policy is ExecutionPolicy.VAST:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
logger.warning(
|
|
"No Vast.ai automatic candidate for component=%s: %s",
|
|
component.representative_node_id,
|
|
exc,
|
|
)
|
|
else:
|
|
environments.append(vast_service.scheduling_state(vast_quote))
|
|
environment_ids = [
|
|
environment.environment_id for environment in environments
|
|
]
|
|
historical_estimates = (
|
|
history.estimates(component_signature, environment_ids)
|
|
if history is not None
|
|
else {}
|
|
)
|
|
component_required_provider = component_required_providers[component_id]
|
|
if preferences.policy is ExecutionPolicy.VAST:
|
|
if component_required_provider is ExecutionProvider.SSH_DOCKER:
|
|
raise ModalPromptValidationError(
|
|
"Vast.ai-only execution cannot run an SSH-only component. "
|
|
"Select Automatic or Self-hosted execution for this workflow."
|
|
)
|
|
component_required_provider = ExecutionProvider.VAST
|
|
requirements = ComponentResourceRequirements(
|
|
minimum_vram_bytes=memory_estimate.minimum_vram_bytes,
|
|
minimum_ram_bytes=memory_estimate.minimum_ram_bytes,
|
|
gpu_required=True,
|
|
architecture="x86_64",
|
|
estimated_execution_seconds=60.0,
|
|
estimated_execution_seconds_by_environment={
|
|
environment_id: estimate.execution_seconds
|
|
for environment_id, estimate in historical_estimates.items()
|
|
},
|
|
preferred_environment_ids=preferences.preferred_environment_ids,
|
|
required_provider=component_required_provider,
|
|
)
|
|
try:
|
|
assignment = _choose_with_idle_ssh_worker_reclaim(
|
|
scheduler=scheduler,
|
|
environments=environments,
|
|
requirements=requirements,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
settings=settings,
|
|
)
|
|
except NoCompatibleExecutionEnvironmentError as exc:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
if assignment.provider is ExecutionProvider.VAST:
|
|
if vast_service is None or vast_quote is None:
|
|
raise ModalPromptValidationError(
|
|
"Vast.ai was selected without a current marketplace quote."
|
|
)
|
|
try:
|
|
if status_callback is None:
|
|
lease = vast_service.acquire_sync(vast_quote)
|
|
else:
|
|
lease = vast_service.acquire_sync(
|
|
vast_quote,
|
|
status_callback=lambda message: status_callback(
|
|
message,
|
|
None,
|
|
None,
|
|
),
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
raise ModalPromptValidationError(
|
|
f"Unable to acquire Vast.ai capacity: {exc}"
|
|
) from exc
|
|
assignment = replace(
|
|
assignment,
|
|
environment_id=lease.environment_id,
|
|
predicted_cost_usd=vast_quote.predicted_incremental_cost_usd,
|
|
reasons=assignment.reasons
|
|
+ (
|
|
f"Vast profile {vast_quote.profile.profile_name}",
|
|
f"idle retention {lease.idle_retention_seconds / 3600:.1f}h",
|
|
),
|
|
)
|
|
assignments[component.representative_node_id] = assignment
|
|
logger.info(
|
|
"Assigned remote component=%s provider=%s environment=%s predicted_cost_usd=%s reasons=%s.",
|
|
component.representative_node_id,
|
|
assignment.provider.value,
|
|
assignment.environment_id,
|
|
assignment.predicted_cost_usd,
|
|
assignment.reasons,
|
|
)
|
|
return assignments
|
|
|
|
|
|
def _plan_component_execution(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
prompt: Mapping[str, Any],
|
|
workflow: Mapping[str, Any] | None,
|
|
settings: ModalSyncSettings,
|
|
status_callback: SetupStatusCallback | None = None,
|
|
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
|
|
plan_callback: ExecutionPlanStatusCallback | None = None,
|
|
occupied_environment_ids: frozenset[str] = frozenset(),
|
|
) -> ComponentExecutionPlan:
|
|
"""Plan through connected configurations or synthesize the legacy behavior."""
|
|
try:
|
|
configuration_set = compile_remote_configuration_set(prompt)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
resolved_llm_profiles = _resolve_prompt_llm_profiles(prompt, settings)
|
|
if configuration_set is None:
|
|
return ComponentExecutionPlan(
|
|
assignments=_plan_component_execution_assignments(
|
|
components=components,
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=settings,
|
|
status_callback=status_callback,
|
|
resolved_llm_profiles=resolved_llm_profiles,
|
|
),
|
|
resolved_llm_profiles=resolved_llm_profiles,
|
|
)
|
|
return _plan_configured_component_execution(
|
|
components=components,
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=settings,
|
|
configuration_set=configuration_set,
|
|
status_callback=status_callback,
|
|
environment_status_callback=environment_status_callback,
|
|
plan_callback=plan_callback,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
resolved_llm_profiles=resolved_llm_profiles,
|
|
)
|
|
|
|
|
|
def _plan_configured_component_execution(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
prompt: Mapping[str, Any],
|
|
workflow: Mapping[str, Any] | None,
|
|
settings: ModalSyncSettings,
|
|
configuration_set: RemoteConfigurationSet,
|
|
status_callback: SetupStatusCallback | None = None,
|
|
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
|
|
plan_callback: ExecutionPlanStatusCallback | None = None,
|
|
occupied_environment_ids: frozenset[str] = frozenset(),
|
|
resolved_llm_profiles: Mapping[str, Any] | None = None,
|
|
) -> ComponentExecutionPlan:
|
|
"""Resolve and prepare a capacity-aware plan from connected configurations."""
|
|
configurations_by_id = {
|
|
configuration.configuration_id: configuration
|
|
for configuration in configuration_set.capacity_configurations
|
|
}
|
|
ssh_hosts_by_id = _probe_configured_ssh_hosts(configuration_set)
|
|
vast_service, vast_unavailable_reason = _configured_vast_service(
|
|
configuration_set,
|
|
settings,
|
|
)
|
|
requirements_by_component = _configured_component_requirements(
|
|
components=components,
|
|
prompt=prompt,
|
|
workflow=workflow,
|
|
settings=settings,
|
|
resolved_llm_profiles=resolved_llm_profiles,
|
|
)
|
|
_prefetch_configured_vast_offers(
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements_by_component,
|
|
vast_service=vast_service,
|
|
)
|
|
environments_by_component, vast_quotes = _configured_candidate_environments(
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements_by_component,
|
|
settings=settings,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_service=vast_service,
|
|
vast_unavailable_reason=vast_unavailable_reason,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
)
|
|
if _reclaim_idle_configured_ssh_capacity(
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements_by_component,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
):
|
|
environments_by_component, vast_quotes = _configured_candidate_environments(
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements_by_component,
|
|
settings=settings,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_service=vast_service,
|
|
vast_unavailable_reason=vast_unavailable_reason,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
)
|
|
_apply_historical_execution_estimates(
|
|
components=components,
|
|
environments_by_component=environments_by_component,
|
|
requirements_by_component=requirements_by_component,
|
|
prompt=prompt,
|
|
settings=settings,
|
|
)
|
|
execution_stages = _component_execution_stages(
|
|
dict(prompt),
|
|
{
|
|
component.representative_node_id: set(component.node_ids)
|
|
for component in components
|
|
},
|
|
)
|
|
try:
|
|
assignments = CostAwareEnvironmentScheduler().plan(
|
|
execution_stages=execution_stages,
|
|
environments_by_component=environments_by_component,
|
|
requirements_by_component=requirements_by_component,
|
|
)
|
|
except NoCompatibleExecutionEnvironmentError as exc:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
|
|
for component_id, assignment in tuple(assignments.items()):
|
|
if assignment.environment_id not in occupied_environment_ids:
|
|
continue
|
|
assignments[component_id] = replace(
|
|
assignment,
|
|
reasons=assignment.reasons
|
|
+ ("capacity available after the earlier workflow completes",),
|
|
)
|
|
|
|
safe_configurations = _safe_remote_configuration_payload(configuration_set)
|
|
if plan_callback is not None:
|
|
plan_callback(
|
|
_planned_execution_assignments_payload(
|
|
assignments,
|
|
components,
|
|
configurations_by_id=configurations_by_id,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_quotes=vast_quotes,
|
|
),
|
|
safe_configurations,
|
|
)
|
|
|
|
vast_leases = _prepare_selected_vast_capacity(
|
|
assignments=assignments,
|
|
configuration_set=configuration_set,
|
|
requirements_by_component=requirements_by_component,
|
|
vast_quotes=vast_quotes,
|
|
vast_service=vast_service,
|
|
status_callback=status_callback,
|
|
environment_status_callback=environment_status_callback,
|
|
)
|
|
if plan_callback is not None and vast_leases:
|
|
plan_callback(
|
|
_planned_execution_assignments_payload(
|
|
assignments,
|
|
components,
|
|
configurations_by_id=configurations_by_id,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_quotes=vast_quotes,
|
|
vast_leases_by_environment=vast_leases,
|
|
),
|
|
safe_configurations,
|
|
)
|
|
return ComponentExecutionPlan(
|
|
assignments=assignments,
|
|
configuration_set=configuration_set,
|
|
configurations_by_id=configurations_by_id,
|
|
safe_configurations=safe_configurations,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_service=vast_service,
|
|
vast_leases_by_environment=vast_leases,
|
|
resolved_llm_profiles=dict(resolved_llm_profiles or {}),
|
|
)
|
|
|
|
|
|
def _probe_configured_ssh_hosts(
|
|
configuration_set: RemoteConfigurationSet,
|
|
) -> dict[str, SshHostConfig]:
|
|
"""Probe all workflow-declared SSH hosts concurrently."""
|
|
ssh_configurations = [
|
|
configuration
|
|
for configuration in configuration_set.configurations
|
|
if isinstance(configuration, SshRemoteConfiguration)
|
|
]
|
|
if not ssh_configurations:
|
|
return {}
|
|
with ThreadPoolExecutor(max_workers=min(8, len(ssh_configurations))) as executor:
|
|
probed_hosts = {
|
|
configuration.configuration_id: future.result()
|
|
for configuration, future in (
|
|
(
|
|
configuration,
|
|
executor.submit(
|
|
_probe_workflow_ssh_configuration,
|
|
configuration,
|
|
),
|
|
)
|
|
for configuration in ssh_configurations
|
|
)
|
|
}
|
|
return probed_hosts
|
|
|
|
|
|
def _configured_vast_service(
|
|
configuration_set: RemoteConfigurationSet,
|
|
settings: ModalSyncSettings,
|
|
) -> tuple[VastService | None, str | None]:
|
|
"""Construct one Vast service only when the connected graph requests it."""
|
|
vast_configured = any(
|
|
isinstance(configuration, VastRemoteConfiguration)
|
|
for configuration in configuration_set.capacity_configurations
|
|
)
|
|
if not vast_configured:
|
|
return None, None
|
|
try:
|
|
return (
|
|
VastService.from_environment(
|
|
settings,
|
|
repo_root=Path(__file__).resolve().parent,
|
|
),
|
|
None,
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
has_fallback_provider = any(
|
|
not isinstance(configuration, VastRemoteConfiguration)
|
|
for configuration in configuration_set.capacity_configurations
|
|
)
|
|
if not has_fallback_provider:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
logger.warning(
|
|
"Skipping unavailable workflow Vast.ai capacity because another "
|
|
"configured provider can be planned: %s",
|
|
exc,
|
|
)
|
|
return None, str(exc)
|
|
|
|
|
|
def _configured_component_requirements(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
prompt: Mapping[str, Any],
|
|
workflow: Mapping[str, Any] | None,
|
|
settings: ModalSyncSettings,
|
|
resolved_llm_profiles: Mapping[str, Any] | None = None,
|
|
) -> dict[str, ComponentResourceRequirements]:
|
|
"""Build provider-neutral requirements for every remote component."""
|
|
preferences = WorkflowExecutionPreferences.from_workflow(workflow)
|
|
active_llm_profiles = (
|
|
dict(resolved_llm_profiles)
|
|
if resolved_llm_profiles is not None
|
|
else _resolve_prompt_llm_profiles(prompt, settings)
|
|
)
|
|
requirements: dict[str, ComponentResourceRequirements] = {}
|
|
for component in components:
|
|
memory_estimate = _component_memory_estimate(
|
|
component,
|
|
prompt,
|
|
preferences,
|
|
settings,
|
|
active_llm_profiles,
|
|
)
|
|
component_id = component.representative_node_id
|
|
requirements[component_id] = ComponentResourceRequirements(
|
|
minimum_vram_bytes=memory_estimate.minimum_vram_bytes,
|
|
minimum_ram_bytes=memory_estimate.minimum_ram_bytes,
|
|
gpu_required=True,
|
|
architecture="x86_64",
|
|
estimated_execution_seconds=60.0,
|
|
preferred_environment_ids=preferences.preferred_environment_ids,
|
|
required_provider=_component_required_provider(
|
|
component,
|
|
prompt,
|
|
active_llm_profiles,
|
|
),
|
|
)
|
|
return requirements
|
|
|
|
|
|
def _configured_candidate_environments(
|
|
*,
|
|
configuration_set: RemoteConfigurationSet,
|
|
requirements_by_component: Mapping[str, ComponentResourceRequirements],
|
|
settings: ModalSyncSettings,
|
|
ssh_hosts_by_id: Mapping[str, SshHostConfig],
|
|
vast_service: VastService | None,
|
|
vast_unavailable_reason: str | None,
|
|
occupied_environment_ids: frozenset[str] = frozenset(),
|
|
) -> tuple[
|
|
dict[str, list[EnvironmentSchedulingState]],
|
|
dict[tuple[str, str], VastProfileQuote],
|
|
]:
|
|
"""Resolve candidate pool states for every component without acquiring capacity."""
|
|
candidates: dict[str, list[EnvironmentSchedulingState]] = {
|
|
component_id: [] for component_id in requirements_by_component
|
|
}
|
|
vast_quotes: dict[tuple[str, str], VastProfileQuote] = {}
|
|
for component_id, requirements in requirements_by_component.items():
|
|
for configuration in configuration_set.capacity_configurations:
|
|
state, quote = _configured_candidate_environment(
|
|
configuration=configuration,
|
|
requirements=requirements,
|
|
settings=settings,
|
|
ssh_hosts_by_id=ssh_hosts_by_id,
|
|
vast_service=vast_service,
|
|
vast_unavailable_reason=vast_unavailable_reason,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
)
|
|
if state is not None:
|
|
candidates[component_id].append(state)
|
|
if quote is not None:
|
|
vast_quotes[(component_id, configuration.configuration_id)] = quote
|
|
return candidates, vast_quotes
|
|
|
|
|
|
def _reclaim_idle_configured_ssh_capacity(
|
|
*,
|
|
configuration_set: RemoteConfigurationSet,
|
|
requirements_by_component: Mapping[str, ComponentResourceRequirements],
|
|
ssh_hosts_by_id: dict[str, SshHostConfig],
|
|
occupied_environment_ids: frozenset[str],
|
|
) -> bool:
|
|
"""Recycle idle configured workers when only their cached memory blocks work."""
|
|
scheduler = CostAwareEnvironmentScheduler()
|
|
reclaimed = False
|
|
for configuration in configuration_set.capacity_configurations:
|
|
if not isinstance(configuration, SshRemoteConfiguration):
|
|
continue
|
|
host = ssh_hosts_by_id[configuration.configuration_id]
|
|
if host.environment_id in occupied_environment_ids:
|
|
continue
|
|
state = replace(
|
|
host.scheduling_state(),
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
)
|
|
optimistic_state = _maximum_capacity_state(state)
|
|
reclaimable = any(
|
|
_optional_scheduler_choice(scheduler, [state], requirements)[0] is None
|
|
and _optional_scheduler_choice(
|
|
scheduler,
|
|
[optimistic_state],
|
|
requirements,
|
|
)[0]
|
|
is not None
|
|
for requirements in requirements_by_component.values()
|
|
)
|
|
if not reclaimable:
|
|
continue
|
|
removed_workers = _remove_idle_ssh_workers_for_reclaim(host)
|
|
if not removed_workers:
|
|
continue
|
|
refreshed_host = _probe_workflow_ssh_configuration(configuration)
|
|
ssh_hosts_by_id[configuration.configuration_id] = refreshed_host
|
|
reclaimed = True
|
|
logger.info(
|
|
"Recycled idle workflow-configured SSH worker(s) environment=%s "
|
|
"containers=%s and refreshed available memory.",
|
|
host.environment_id,
|
|
removed_workers,
|
|
)
|
|
return reclaimed
|
|
|
|
|
|
def _prefetch_configured_vast_offers(
|
|
*,
|
|
configuration_set: RemoteConfigurationSet,
|
|
requirements_by_component: Mapping[str, ComponentResourceRequirements],
|
|
vast_service: VastService | None,
|
|
) -> None:
|
|
"""Warm distinct Vast searches in parallel before deterministic planning."""
|
|
if vast_service is None:
|
|
return
|
|
profiles = [
|
|
configuration.profile
|
|
for configuration in configuration_set.configurations
|
|
if isinstance(configuration, VastRemoteConfiguration)
|
|
]
|
|
prefetch = getattr(vast_service, "prefetch_offers_sync", None)
|
|
if not callable(prefetch):
|
|
return
|
|
try:
|
|
prefetch(
|
|
profiles,
|
|
tuple(
|
|
VastSearchRequirements(
|
|
minimum_vram_bytes=requirements.minimum_vram_bytes,
|
|
minimum_ram_bytes=requirements.minimum_ram_bytes,
|
|
)
|
|
for requirements in requirements_by_component.values()
|
|
),
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Unable to prefetch configured Vast.ai marketplace offers: %s",
|
|
exc,
|
|
)
|
|
|
|
|
|
def _configured_candidate_environment(
|
|
*,
|
|
configuration: RemoteConfiguration,
|
|
requirements: ComponentResourceRequirements,
|
|
settings: ModalSyncSettings,
|
|
ssh_hosts_by_id: Mapping[str, SshHostConfig],
|
|
vast_service: VastService | None,
|
|
vast_unavailable_reason: str | None,
|
|
occupied_environment_ids: frozenset[str] = frozenset(),
|
|
) -> tuple[EnvironmentSchedulingState | None, VastProfileQuote | None]:
|
|
"""Resolve one configuration as a candidate for one component."""
|
|
if isinstance(configuration, ModalRemoteConfiguration):
|
|
return _modal_configuration_environment_state(configuration, settings), None
|
|
if isinstance(configuration, SshRemoteConfiguration):
|
|
host = ssh_hosts_by_id[configuration.configuration_id]
|
|
state = replace(
|
|
host.scheduling_state(),
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
)
|
|
if state.environment_id in occupied_environment_ids:
|
|
state = _maximum_capacity_state(state)
|
|
logger.info(
|
|
"Planning queued work against reclaimable SSH capacity "
|
|
"environment=%s because an earlier workflow owns that host.",
|
|
state.environment_id,
|
|
)
|
|
return state, None
|
|
if not isinstance(configuration, VastRemoteConfiguration):
|
|
raise TypeError(
|
|
f"Unsupported remote configuration type {type(configuration).__name__}."
|
|
)
|
|
if vast_service is None:
|
|
return EnvironmentSchedulingState(
|
|
environment_id=f"vast:{configuration.configuration_id}",
|
|
provider=ExecutionProvider.VAST,
|
|
enabled=True,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
cost_usd_per_second=None,
|
|
capabilities=None,
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
unavailable_reason=(
|
|
vast_unavailable_reason
|
|
or "Vast.ai controller is unavailable."
|
|
),
|
|
maximum_workers=configuration.capacity_limit,
|
|
), None
|
|
try:
|
|
quote = vast_service.quote_best_profile_sync(
|
|
[configuration.profile],
|
|
minimum_vram_bytes=requirements.minimum_vram_bytes,
|
|
minimum_ram_bytes=requirements.minimum_ram_bytes,
|
|
predicted_execution_seconds=requirements.estimated_execution_seconds,
|
|
)
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Vast configuration has no candidate configuration=%s: %s",
|
|
configuration.display_name,
|
|
exc,
|
|
)
|
|
return EnvironmentSchedulingState(
|
|
environment_id=f"vast:{configuration.configuration_id}",
|
|
provider=ExecutionProvider.VAST,
|
|
enabled=True,
|
|
health=EnvironmentHealth.UNAVAILABLE,
|
|
cost_usd_per_second=None,
|
|
capabilities=None,
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
unavailable_reason=str(exc),
|
|
maximum_workers=configuration.capacity_limit,
|
|
), None
|
|
state = vast_service.scheduling_state(quote)
|
|
return replace(
|
|
state,
|
|
environment_id=f"vast:{configuration.configuration_id}",
|
|
configuration_id=configuration.configuration_id,
|
|
display_name=configuration.display_name,
|
|
maximum_workers=configuration.capacity_limit,
|
|
), quote
|
|
|
|
|
|
def _apply_historical_execution_estimates(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
environments_by_component: Mapping[str, list[EnvironmentSchedulingState]],
|
|
requirements_by_component: dict[str, ComponentResourceRequirements],
|
|
prompt: Mapping[str, Any],
|
|
settings: ModalSyncSettings,
|
|
) -> None:
|
|
"""Add environment-specific historical timing estimates in place."""
|
|
history = _execution_history(settings)
|
|
if history is None:
|
|
return
|
|
for component in components:
|
|
component_id = component.representative_node_id
|
|
environment_ids = [
|
|
environment.environment_id
|
|
for environment in environments_by_component.get(component_id, [])
|
|
]
|
|
estimates = history.estimates(
|
|
_component_execution_signature(component, prompt),
|
|
environment_ids,
|
|
)
|
|
requirements_by_component[component_id] = replace(
|
|
requirements_by_component[component_id],
|
|
estimated_execution_seconds_by_environment={
|
|
environment_id: estimate.execution_seconds
|
|
for environment_id, estimate in estimates.items()
|
|
},
|
|
)
|
|
|
|
|
|
def _prepare_selected_vast_capacity(
|
|
*,
|
|
assignments: dict[str, ExecutionAssignment],
|
|
configuration_set: RemoteConfigurationSet,
|
|
requirements_by_component: Mapping[str, ComponentResourceRequirements],
|
|
vast_quotes: Mapping[tuple[str, str], VastProfileQuote],
|
|
vast_service: VastService | None,
|
|
status_callback: SetupStatusCallback | None = None,
|
|
environment_status_callback: EnvironmentSetupStatusCallback | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Acquire only Vast slots selected by the completed provider-neutral plan."""
|
|
selected_slots: dict[tuple[str, int], list[str]] = defaultdict(list)
|
|
for component_id, assignment in assignments.items():
|
|
if assignment.provider is not ExecutionProvider.VAST:
|
|
continue
|
|
if assignment.configuration_id is None:
|
|
raise ModalPromptValidationError(
|
|
"Vast assignment is missing its configuration identity."
|
|
)
|
|
selected_slots[
|
|
(assignment.configuration_id, assignment.capacity_slot_index)
|
|
].append(component_id)
|
|
if not selected_slots:
|
|
return {}
|
|
if vast_service is None:
|
|
raise ModalPromptValidationError(
|
|
"Vast capacity was selected without an initialized Vast service."
|
|
)
|
|
|
|
configurations_by_id = {
|
|
configuration.configuration_id: configuration
|
|
for configuration in configuration_set.configurations
|
|
if isinstance(configuration, VastRemoteConfiguration)
|
|
}
|
|
leases_by_environment: dict[str, Any] = {}
|
|
ordered_slots = sorted(selected_slots.items())
|
|
total_slots = len(ordered_slots)
|
|
for slot_number, ((configuration_id, slot_index), component_ids) in enumerate(
|
|
ordered_slots,
|
|
start=1,
|
|
):
|
|
configuration = configurations_by_id[configuration_id]
|
|
planned_environment_id = assignments[component_ids[0]].environment_id
|
|
quote = _quote_selected_vast_slot(
|
|
configuration=configuration,
|
|
component_ids=component_ids,
|
|
requirements_by_component=requirements_by_component,
|
|
vast_quotes=vast_quotes,
|
|
vast_service=vast_service,
|
|
)
|
|
try:
|
|
_emit_vast_capacity_status(
|
|
environment_id=planned_environment_id,
|
|
message=f"Acquiring Vast.ai capacity {slot_number} of {total_slots}",
|
|
current=slot_number - 1,
|
|
total=total_slots,
|
|
status_callback=status_callback,
|
|
environment_status_callback=environment_status_callback,
|
|
)
|
|
if status_callback is None and environment_status_callback is None:
|
|
lease = vast_service.acquire_sync(quote, slot=slot_index)
|
|
else:
|
|
lease = vast_service.acquire_sync(
|
|
quote,
|
|
slot=slot_index,
|
|
status_callback=(
|
|
lambda message, environment_id=planned_environment_id,
|
|
current=slot_number - 1: _emit_vast_capacity_status(
|
|
environment_id=environment_id,
|
|
message=message,
|
|
current=current,
|
|
total=total_slots,
|
|
status_callback=status_callback,
|
|
environment_status_callback=environment_status_callback,
|
|
)
|
|
),
|
|
)
|
|
except SyncCancelledError:
|
|
raise
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
raise ModalPromptValidationError(
|
|
f"Unable to acquire Vast.ai capacity for configuration "
|
|
f"{configuration.display_name!r} slot {slot_index}: {exc}"
|
|
) from exc
|
|
leases_by_environment[lease.environment_id] = lease
|
|
cost_share = quote.predicted_incremental_cost_usd / len(component_ids)
|
|
for component_id in component_ids:
|
|
assignment = assignments[component_id]
|
|
assignments[component_id] = replace(
|
|
assignment,
|
|
environment_id=lease.environment_id,
|
|
predicted_cost_usd=cost_share,
|
|
reasons=assignment.reasons
|
|
+ (
|
|
f"Vast profile {configuration.display_name}",
|
|
f"Vast capacity slot {slot_index}",
|
|
),
|
|
)
|
|
_emit_vast_capacity_status(
|
|
environment_id=lease.environment_id,
|
|
message=f"Vast.ai capacity {slot_number} of {total_slots} is ready",
|
|
environment_message="Vast.ai worker ready; preparing remote assets next",
|
|
current=slot_number,
|
|
total=total_slots,
|
|
status_callback=status_callback,
|
|
environment_status_callback=environment_status_callback,
|
|
)
|
|
return leases_by_environment
|
|
|
|
|
|
def _emit_vast_capacity_status(
|
|
*,
|
|
environment_id: str,
|
|
message: str,
|
|
environment_message: str | None = None,
|
|
current: int,
|
|
total: int,
|
|
status_callback: SetupStatusCallback | None,
|
|
environment_status_callback: EnvironmentSetupStatusCallback | None,
|
|
) -> None:
|
|
"""Publish one Vast setup update at prompt and environment scopes."""
|
|
if status_callback is not None:
|
|
status_callback(message, current, total)
|
|
if environment_status_callback is not None:
|
|
environment_status_callback(
|
|
environment_id,
|
|
environment_message or message,
|
|
None,
|
|
None,
|
|
)
|
|
|
|
|
|
def _quote_selected_vast_slot(
|
|
*,
|
|
configuration: VastRemoteConfiguration,
|
|
component_ids: list[str],
|
|
requirements_by_component: Mapping[str, ComponentResourceRequirements],
|
|
vast_quotes: Mapping[tuple[str, str], VastProfileQuote],
|
|
vast_service: VastService,
|
|
) -> VastProfileQuote:
|
|
"""Quote aggregate slot requirements immediately before acquisition."""
|
|
minimum_vram_bytes = max(
|
|
requirements_by_component[component_id].minimum_vram_bytes
|
|
for component_id in component_ids
|
|
)
|
|
minimum_ram_bytes = max(
|
|
requirements_by_component[component_id].minimum_ram_bytes
|
|
for component_id in component_ids
|
|
)
|
|
predicted_execution_seconds = sum(
|
|
requirements_by_component[component_id].estimated_execution_seconds
|
|
for component_id in component_ids
|
|
)
|
|
if len(component_ids) == 1:
|
|
existing_quote = vast_quotes.get(
|
|
(component_ids[0], configuration.configuration_id)
|
|
)
|
|
if existing_quote is not None:
|
|
return existing_quote
|
|
return vast_service.quote_best_profile_sync(
|
|
[configuration.profile],
|
|
minimum_vram_bytes=minimum_vram_bytes,
|
|
minimum_ram_bytes=minimum_ram_bytes,
|
|
predicted_execution_seconds=predicted_execution_seconds,
|
|
)
|