Add conservative remote memory admission
This commit is contained in:
@@ -113,6 +113,8 @@ Workflow policy is stored under `workflow.extra.remote_execution` and travels wi
|
||||
- **Automatic** considers ready SSH hosts and the selected Modal GPU, rejects hard incompatibilities such as insufficient VRAM, then prefers known lowest predicted cost with deterministic tie-breaking.
|
||||
- **Automatically place eligible nodes** makes the planner select the executable prompt without requiring each node toggle. Leave it off to keep manual component boundaries.
|
||||
|
||||
Before comparing cost, the planner resolves model files referenced by each remote component and derives conservative memory floors. It budgets the largest primary model that may be resident, adds LoRA/adapter/ControlNet weights, applies 20% weight overhead, and reserves another 4 GiB for activations and runtime state. RAM receives the resident weights plus 4 GiB. Hosts are checked against currently free VRAM and available RAM when the capability probe reports them, falling back to total capacity only when live availability is unknown. An undersized zero-cost host is therefore rejected before any model upload; automatic policy can fall back to a compatible Modal GPU, while self-hosted-only policy returns the per-host memory shortfall.
|
||||
|
||||
Images are immutable and named from the local node-pack source, selected ComfyUI source, custom-node requirements, and runtime-shaping settings. Warm containers use a persistent named volume mounted at `/storage`; assets and custom-node bundles are content-addressed, and a remote volume epoch prevents a deleted/recreated volume from reusing stale local index records. Requests and tensor payloads travel as bounded, length-prefixed binary frames through `ssh docker exec`; no daemon TCP port, worker port, or inbound firewall rule is required. Generated ComfyUI files are bundled into the result, integrity-checked, and atomically restored beneath the local output directory just as they are for Modal.
|
||||
|
||||
Controller SSH sessions disable forwards inherited from the selected alias, so existing development tunnels cannot collide with worker-management connections. Image builds explicitly load their result into the remote daemon, including when its selected BuildKit builder uses the isolated `docker-container` driver. The worker entrypoint keeps Modal-only package initialization lazy, allowing the same source tree to run as a top-level OCI application while ordinary ComfyUI imports retain package-relative behavior. Custom-node requirements are installed into the image, except packages owned by the immutable Python/CUDA runtime—such as NumPy, OpenCV, Pillow, PyTorch, Triton, and vLLM—whose custom-node constraints are ignored before the final supported pins are applied. A custom node that is genuinely incompatible with those supported versions will fail explicitly during import rather than silently replacing the worker stack.
|
||||
|
||||
+162
-6
@@ -7,6 +7,7 @@ import copy
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -58,10 +59,12 @@ from .execution_environments import (
|
||||
from .execution_history import ExecutionHistory
|
||||
from .sync_engine import (
|
||||
AssetSyncRequestCache,
|
||||
MODEL_FILE_EXTENSIONS,
|
||||
ModalAssetSyncEngine,
|
||||
ModalVolumeBackend,
|
||||
SyncedAsset,
|
||||
modal,
|
||||
resolve_model_path,
|
||||
)
|
||||
from .remote_hosts import RemoteExecutionConfig, RemoteHostRegistry, SshHostConfig
|
||||
from .ssh_docker import SshDockerController, SshDockerVolumeBackend
|
||||
@@ -132,6 +135,25 @@ _SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset(
|
||||
"snapshot_profile_key",
|
||||
}
|
||||
)
|
||||
_MODEL_VRAM_WEIGHT_MULTIPLIER = 1.20
|
||||
_MODEL_VRAM_HEADROOM_BYTES = 4 * 1024**3
|
||||
_MODEL_RAM_HEADROOM_BYTES = 4 * 1024**3
|
||||
_ADDITIVE_MODEL_CLASS_TOKENS = (
|
||||
"adapter",
|
||||
"controlnet",
|
||||
"ipadapter",
|
||||
"lora",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ComponentMemoryEstimate:
|
||||
"""Describe conservative scheduler memory floors inferred from model assets."""
|
||||
|
||||
minimum_vram_bytes: int
|
||||
minimum_ram_bytes: int
|
||||
model_asset_count: int = 0
|
||||
largest_model_bytes: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -453,8 +475,23 @@ def _component_minimum_vram_bytes(
|
||||
component: RemoteComponentPlan,
|
||||
prompt: Mapping[str, Any],
|
||||
preferences: WorkflowExecutionPreferences,
|
||||
settings: ModalSyncSettings,
|
||||
) -> int:
|
||||
"""Estimate a conservative component VRAM floor from explicit LLM profiles."""
|
||||
"""Return the conservative GPU VRAM floor for one remote component."""
|
||||
return _component_memory_estimate(
|
||||
component,
|
||||
prompt,
|
||||
preferences,
|
||||
settings,
|
||||
).minimum_vram_bytes
|
||||
|
||||
|
||||
def _component_profile_minimum_vram_bytes(
|
||||
component: RemoteComponentPlan,
|
||||
prompt: Mapping[str, Any],
|
||||
preferences: WorkflowExecutionPreferences,
|
||||
) -> int:
|
||||
"""Estimate a component VRAM floor from workflow and LLM profile metadata."""
|
||||
minimum_vram_bytes = preferences.minimum_vram_bytes
|
||||
try:
|
||||
from .llm_profiles import get_llm_profile
|
||||
@@ -483,6 +520,111 @@ def _component_minimum_vram_bytes(
|
||||
return minimum_vram_bytes
|
||||
|
||||
|
||||
def _iter_prompt_string_values(value: object) -> Iterator[str]:
|
||||
"""Yield nested prompt string values while ignoring graph links and scalars."""
|
||||
if isinstance(value, str):
|
||||
yield value
|
||||
return
|
||||
if isinstance(value, Mapping):
|
||||
for child in value.values():
|
||||
yield from _iter_prompt_string_values(child)
|
||||
return
|
||||
if isinstance(value, (list, tuple)) and not _is_link(value):
|
||||
for child in value:
|
||||
yield from _iter_prompt_string_values(child)
|
||||
|
||||
|
||||
def _is_additive_model_node(class_type: str) -> bool:
|
||||
"""Return whether a loader contributes weights alongside a primary model."""
|
||||
normalized = class_type.casefold()
|
||||
return any(token in normalized for token in _ADDITIVE_MODEL_CLASS_TOKENS)
|
||||
|
||||
|
||||
def _component_model_asset_sizes(
|
||||
component: RemoteComponentPlan,
|
||||
prompt: Mapping[str, Any],
|
||||
settings: ModalSyncSettings,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
"""Return unique primary and additive model sizes referenced by a component."""
|
||||
asset_roles: dict[Path, bool] = {}
|
||||
for node_id in component.node_ids:
|
||||
prompt_node = prompt.get(node_id)
|
||||
if not isinstance(prompt_node, Mapping):
|
||||
continue
|
||||
inputs = prompt_node.get("inputs")
|
||||
if not isinstance(inputs, Mapping):
|
||||
continue
|
||||
additive = _is_additive_model_node(str(prompt_node.get("class_type") or ""))
|
||||
for value in _iter_prompt_string_values(inputs):
|
||||
resolved_path = resolve_model_path(
|
||||
value,
|
||||
comfyui_root=getattr(settings, "comfyui_root", None),
|
||||
extensions=MODEL_FILE_EXTENSIONS,
|
||||
)
|
||||
if resolved_path is None:
|
||||
continue
|
||||
asset_roles[resolved_path] = (
|
||||
asset_roles.get(resolved_path, True) and additive
|
||||
)
|
||||
|
||||
primary_sizes: list[int] = []
|
||||
additive_sizes: list[int] = []
|
||||
for asset_path, additive in asset_roles.items():
|
||||
try:
|
||||
size_bytes = asset_path.stat().st_size
|
||||
except OSError as exc:
|
||||
logger.warning(
|
||||
"Unable to inspect model size for placement path=%s: %s",
|
||||
asset_path,
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
if size_bytes <= 0:
|
||||
continue
|
||||
target = additive_sizes if additive else primary_sizes
|
||||
target.append(size_bytes)
|
||||
return tuple(primary_sizes), tuple(additive_sizes)
|
||||
|
||||
|
||||
def _component_memory_estimate(
|
||||
component: RemoteComponentPlan,
|
||||
prompt: Mapping[str, Any],
|
||||
preferences: WorkflowExecutionPreferences,
|
||||
settings: ModalSyncSettings,
|
||||
) -> ComponentMemoryEstimate:
|
||||
"""Infer conservative RAM and VRAM floors from resident model weight sizes."""
|
||||
profile_vram_bytes = _component_profile_minimum_vram_bytes(
|
||||
component,
|
||||
prompt,
|
||||
preferences,
|
||||
)
|
||||
primary_sizes, additive_sizes = _component_model_asset_sizes(
|
||||
component,
|
||||
prompt,
|
||||
settings,
|
||||
)
|
||||
all_sizes = primary_sizes + additive_sizes
|
||||
if not all_sizes:
|
||||
return ComponentMemoryEstimate(
|
||||
minimum_vram_bytes=profile_vram_bytes,
|
||||
minimum_ram_bytes=0,
|
||||
)
|
||||
|
||||
primary_peak_bytes = max(primary_sizes or all_sizes)
|
||||
additive_bytes = sum(additive_sizes) if primary_sizes else 0
|
||||
resident_weight_bytes = primary_peak_bytes + additive_bytes
|
||||
model_vram_bytes = (
|
||||
math.ceil(resident_weight_bytes * _MODEL_VRAM_WEIGHT_MULTIPLIER)
|
||||
+ _MODEL_VRAM_HEADROOM_BYTES
|
||||
)
|
||||
return ComponentMemoryEstimate(
|
||||
minimum_vram_bytes=max(profile_vram_bytes, model_vram_bytes),
|
||||
minimum_ram_bytes=resident_weight_bytes + _MODEL_RAM_HEADROOM_BYTES,
|
||||
model_asset_count=len(all_sizes),
|
||||
largest_model_bytes=max(all_sizes),
|
||||
)
|
||||
|
||||
|
||||
def _component_execution_signature(
|
||||
component: RemoteComponentPlan,
|
||||
prompt: Mapping[str, Any],
|
||||
@@ -574,12 +716,25 @@ def _plan_component_execution_assignments(
|
||||
if history is not None
|
||||
else {}
|
||||
)
|
||||
memory_estimate = _component_memory_estimate(
|
||||
component,
|
||||
prompt,
|
||||
preferences,
|
||||
settings,
|
||||
)
|
||||
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.representative_node_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,
|
||||
)
|
||||
requirements = ComponentResourceRequirements(
|
||||
minimum_vram_bytes=_component_minimum_vram_bytes(
|
||||
component,
|
||||
prompt,
|
||||
preferences,
|
||||
),
|
||||
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,
|
||||
@@ -5127,6 +5282,7 @@ def rewrite_prompt_for_modal(
|
||||
component,
|
||||
rewritten_prompt,
|
||||
preferences,
|
||||
resolved_settings,
|
||||
),
|
||||
)
|
||||
anchor_assignment = assignments_by_component_id[
|
||||
|
||||
@@ -467,3 +467,4 @@
|
||||
- Reapplying deployment-owned pins after custom-node requirements is too late when a conflicting requirement has no wheel for the worker Python version: pip can fail while attempting a source build before the repair layer runs. Filter requirements for packages owned by the immutable Python/CUDA stack, retain all other custom-node dependencies in the fingerprint, and let import validation expose genuine incompatibility.
|
||||
- A selected BuildKit `docker-container` driver does not automatically place a successful `docker build` result in the host image store. Pass `--load` so the fingerprinted image is available to the subsequent `docker run`; otherwise the build succeeds into cache and Docker tries to pull the private local tag from a registry.
|
||||
- Running `python -m remote.ssh_worker` imports `remote/__init__.py` before the worker module. Eagerly importing `remote.modal_app` there breaks because its parent-relative imports assume the node pack has a parent package, while the OCI source root deliberately exposes `remote` at top level. Keep those convenience exports lazy and regression-test the exact subprocess entrypoint used by the container.
|
||||
- Cost optimization is safe only after component resource admission. Ordinary loader nodes often have no explicit VRAM declaration, so resolve their local model files before placement, budget the largest primary weight set plus additive LoRA/adapter/ControlNet weights, apply runtime headroom, and compare against probed free VRAM and available RAM. Otherwise an explicit zero-cost 24 GiB host can beat a large Modal GPU for a 66 GiB BF16 model and begin a doomed upload before failing at execution.
|
||||
|
||||
@@ -112,6 +112,25 @@ class EnvironmentCapabilities:
|
||||
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 {
|
||||
@@ -322,6 +341,13 @@ class CostAwareEnvironmentScheduler:
|
||||
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
|
||||
@@ -370,15 +396,28 @@ class CostAwareEnvironmentScheduler:
|
||||
f"architecture {capabilities.architecture!r} does not satisfy "
|
||||
f"{requirements.architecture!r}"
|
||||
)
|
||||
if capabilities.total_ram_bytes < requirements.minimum_ram_bytes:
|
||||
return "insufficient RAM"
|
||||
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_vram_bytes < requirements.minimum_vram_bytes:
|
||||
return "insufficient GPU VRAM"
|
||||
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(
|
||||
|
||||
+56
-50
@@ -21,7 +21,20 @@ from .settings import ModalSyncSettings, get_settings
|
||||
logger = logging.getLogger(__name__)
|
||||
SyncStatusCallback = Callable[[str, int | None, int | None], None]
|
||||
|
||||
_SYNC_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".vae"}
|
||||
_SYNC_EXTENSIONS = frozenset({".safetensors", ".ckpt", ".pt", ".vae"})
|
||||
MODEL_FILE_EXTENSIONS = frozenset(
|
||||
{
|
||||
".bin",
|
||||
".ckpt",
|
||||
".engine",
|
||||
".gguf",
|
||||
".onnx",
|
||||
".pt",
|
||||
".pth",
|
||||
".safetensors",
|
||||
".vae",
|
||||
}
|
||||
)
|
||||
_SKIP_DIRS = {
|
||||
".git",
|
||||
".mypy_cache",
|
||||
@@ -43,19 +56,7 @@ _SKIP_FILE_SUFFIXES = {
|
||||
".tmp",
|
||||
}
|
||||
_CUSTOM_NODES_MANIFEST_VERSION = 2
|
||||
_CUSTOM_NODE_ASSET_SUFFIXES = frozenset(
|
||||
{
|
||||
".bin",
|
||||
".ckpt",
|
||||
".engine",
|
||||
".gguf",
|
||||
".onnx",
|
||||
".pt",
|
||||
".pth",
|
||||
".safetensors",
|
||||
".vae",
|
||||
}
|
||||
)
|
||||
_CUSTOM_NODE_ASSET_SUFFIXES = MODEL_FILE_EXTENSIONS
|
||||
|
||||
try:
|
||||
import modal # type: ignore
|
||||
@@ -63,6 +64,39 @@ except ModuleNotFoundError: # pragma: no cover - exercised when Modal SDK is un
|
||||
modal = None
|
||||
|
||||
|
||||
def resolve_model_path(
|
||||
value: str,
|
||||
*,
|
||||
comfyui_root: Path | None = None,
|
||||
extensions: frozenset[str] = _SYNC_EXTENSIONS,
|
||||
) -> Path | None:
|
||||
"""Resolve one prompt string to a local model file when possible."""
|
||||
path = Path(value).expanduser()
|
||||
if path.suffix.lower() not in extensions:
|
||||
return None
|
||||
if path.is_file():
|
||||
return path.resolve()
|
||||
if os.path.isabs(value):
|
||||
return None
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
except ModuleNotFoundError:
|
||||
folder_paths = None
|
||||
|
||||
if folder_paths is not None:
|
||||
for folder_name in folder_paths.folder_names_and_paths:
|
||||
full_path = folder_paths.get_full_path(folder_name, value)
|
||||
if full_path is not None:
|
||||
return Path(full_path).resolve()
|
||||
|
||||
if comfyui_root is not None:
|
||||
candidate = comfyui_root / value
|
||||
if candidate.is_file():
|
||||
return candidate.resolve()
|
||||
return None
|
||||
|
||||
|
||||
def _emit_sync_status(
|
||||
status_callback: SyncStatusCallback | None,
|
||||
message: str,
|
||||
@@ -1082,42 +1116,14 @@ class ModalAssetSyncEngine:
|
||||
if value in self._path_resolution_cache:
|
||||
cached = self._path_resolution_cache[value]
|
||||
return Path(cached) if cached is not None else None
|
||||
|
||||
path = Path(value).expanduser()
|
||||
if path.suffix.lower() not in _SYNC_EXTENSIONS:
|
||||
self._path_resolution_cache[value] = None
|
||||
return None
|
||||
if path.is_file():
|
||||
resolved = path.resolve()
|
||||
self._path_resolution_cache[value] = str(resolved)
|
||||
return resolved
|
||||
|
||||
if os.path.isabs(value):
|
||||
self._path_resolution_cache[value] = None
|
||||
return None
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
except ModuleNotFoundError:
|
||||
folder_paths = None
|
||||
|
||||
if folder_paths is not None:
|
||||
for folder_name in folder_paths.folder_names_and_paths:
|
||||
full_path = folder_paths.get_full_path(folder_name, value)
|
||||
if full_path is not None:
|
||||
resolved = Path(full_path).resolve()
|
||||
self._path_resolution_cache[value] = str(resolved)
|
||||
return resolved
|
||||
|
||||
if self.settings.comfyui_root is not None:
|
||||
candidate = self.settings.comfyui_root / value
|
||||
if candidate.is_file():
|
||||
resolved = candidate.resolve()
|
||||
self._path_resolution_cache[value] = str(resolved)
|
||||
return resolved
|
||||
|
||||
self._path_resolution_cache[value] = None
|
||||
return None
|
||||
resolved = resolve_model_path(
|
||||
value,
|
||||
comfyui_root=self.settings.comfyui_root,
|
||||
)
|
||||
self._path_resolution_cache[value] = (
|
||||
str(resolved) if resolved is not None else None
|
||||
)
|
||||
return resolved
|
||||
|
||||
def _collect_syncable_asset_paths(self, value: Any) -> list[Path]:
|
||||
"""Return all prompt asset paths that resolve to syncable local files."""
|
||||
|
||||
@@ -449,7 +449,11 @@ def test_automatic_policy_assigns_component_to_lower_cost_ready_host(
|
||||
"_schedulable_ssh_hosts",
|
||||
lambda _settings: (host,),
|
||||
)
|
||||
monkeypatch.setattr(api_intercept_module, "_execution_history", lambda _settings: None)
|
||||
monkeypatch.setattr(
|
||||
api_intercept_module,
|
||||
"_execution_history",
|
||||
lambda _settings: None,
|
||||
)
|
||||
|
||||
assignments = api_intercept_module._plan_component_execution_assignments(
|
||||
components=[component],
|
||||
@@ -469,6 +473,100 @@ def test_automatic_policy_assigns_component_to_lower_cost_ready_host(
|
||||
assert assignments["1"].environment_id == "cheap-host"
|
||||
|
||||
|
||||
def test_automatic_policy_rejects_zero_cost_host_for_oversized_model(
|
||||
api_intercept_module: Any,
|
||||
remote_hosts_module: Any,
|
||||
execution_environments_module: Any,
|
||||
monkeypatch: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Model weights plus headroom must exclude a cheap host before cost ranking."""
|
||||
module = execution_environments_module
|
||||
model_path = tmp_path / "minimax_h3_bf16.safetensors"
|
||||
with model_path.open("wb") as model_file:
|
||||
model_file.truncate(66 * 1024**3)
|
||||
|
||||
capabilities = module.EnvironmentCapabilities(
|
||||
architecture="x86_64",
|
||||
operating_system="linux",
|
||||
cpu_count=16,
|
||||
total_ram_bytes=64 * 1024**3,
|
||||
available_ram_bytes=60 * 1024**3,
|
||||
available_disk_bytes=1024**4,
|
||||
docker_version="28.0.0",
|
||||
docker_rootless=False,
|
||||
nvidia_container_runtime=True,
|
||||
gpus=(module.GpuCapability("GPU-4090", "RTX 4090", 24 * 1024**3),),
|
||||
probed_at_epoch=2_000_000_000.0,
|
||||
)
|
||||
host = remote_hosts_module.SshHostConfig(
|
||||
environment_id="lambda",
|
||||
display_name="Lambda 4090",
|
||||
ssh_target="lambda",
|
||||
cost_usd_per_second=0.0,
|
||||
capabilities=capabilities,
|
||||
health=module.EnvironmentHealth.READY,
|
||||
)
|
||||
component = api_intercept_module.RemoteComponentPlan(
|
||||
node_ids=["6"],
|
||||
representative_node_id="6",
|
||||
boundary_inputs=[],
|
||||
boundary_outputs=[],
|
||||
execute_node_ids=["6"],
|
||||
contains_output_node=False,
|
||||
)
|
||||
prompt = {
|
||||
"6": {
|
||||
"class_type": "UNETLoader",
|
||||
"inputs": {"unet_name": str(model_path), "weight_dtype": "default"},
|
||||
}
|
||||
}
|
||||
workflow = {
|
||||
"extra": {
|
||||
"remote_execution": {
|
||||
"policy": "automatic",
|
||||
"auto_place": True,
|
||||
}
|
||||
}
|
||||
}
|
||||
settings = SimpleNamespace(
|
||||
modal_gpu="H200",
|
||||
max_containers=1,
|
||||
comfyui_root=None,
|
||||
)
|
||||
preferences = module.WorkflowExecutionPreferences.from_workflow(workflow)
|
||||
estimate = api_intercept_module._component_memory_estimate(
|
||||
component,
|
||||
prompt,
|
||||
preferences,
|
||||
settings,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
api_intercept_module,
|
||||
"_schedulable_ssh_hosts",
|
||||
lambda _settings: (host,),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
api_intercept_module,
|
||||
"_execution_history",
|
||||
lambda _settings: None,
|
||||
)
|
||||
|
||||
assignments = api_intercept_module._plan_component_execution_assignments(
|
||||
components=[component],
|
||||
prompt=prompt,
|
||||
workflow=workflow,
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert estimate.model_asset_count == 1
|
||||
assert estimate.largest_model_bytes == 66 * 1024**3
|
||||
assert 24 * 1024**3 < estimate.minimum_vram_bytes < 96 * 1024**3
|
||||
assert estimate.minimum_ram_bytes == 70 * 1024**3
|
||||
assert assignments["6"].provider is module.ExecutionProvider.MODAL
|
||||
assert assignments["6"].environment_id == "modal:H200"
|
||||
|
||||
|
||||
def test_planner_keeps_unmarked_llm_local_between_remote_text_nodes(
|
||||
api_intercept_module: Any,
|
||||
) -> None:
|
||||
|
||||
@@ -7,7 +7,13 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
|
||||
def _capabilities(module: Any, *, vram_gb: int, ram_gb: int = 64) -> Any:
|
||||
def _capabilities(
|
||||
module: Any,
|
||||
*,
|
||||
vram_gb: int,
|
||||
ram_gb: int = 64,
|
||||
free_vram_gb: int | None = None,
|
||||
) -> Any:
|
||||
"""Return deterministic Linux GPU capabilities for scheduler tests."""
|
||||
return module.EnvironmentCapabilities(
|
||||
architecture="x86_64",
|
||||
@@ -24,6 +30,11 @@ def _capabilities(module: Any, *, vram_gb: int, ram_gb: int = 64) -> Any:
|
||||
uuid=f"GPU-{vram_gb}",
|
||||
name=f"Test GPU {vram_gb}",
|
||||
total_vram_bytes=vram_gb * 1024**3,
|
||||
free_vram_bytes=(
|
||||
free_vram_gb * 1024**3
|
||||
if free_vram_gb is not None
|
||||
else None
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -81,6 +92,34 @@ def test_scheduler_rejects_hosts_without_required_vram(
|
||||
)
|
||||
|
||||
|
||||
def test_scheduler_uses_probed_free_vram_for_admission(
|
||||
execution_environments_module: Any,
|
||||
) -> None:
|
||||
"""A busy GPU must not be admitted based only on nameplate VRAM."""
|
||||
module = execution_environments_module
|
||||
busy_environment = module.EnvironmentSchedulingState(
|
||||
environment_id="busy",
|
||||
provider=module.ExecutionProvider.SSH_DOCKER,
|
||||
enabled=True,
|
||||
health=module.EnvironmentHealth.READY,
|
||||
cost_usd_per_second=0.0,
|
||||
capabilities=_capabilities(
|
||||
module,
|
||||
vram_gb=80,
|
||||
free_vram_gb=16,
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
module.NoCompatibleExecutionEnvironmentError,
|
||||
match=r"16\.00 GiB available, 40\.00 GiB required",
|
||||
):
|
||||
module.CostAwareEnvironmentScheduler().choose(
|
||||
[busy_environment],
|
||||
module.ComponentResourceRequirements(minimum_vram_bytes=40 * 1024**3),
|
||||
)
|
||||
|
||||
|
||||
def test_scheduler_places_known_cost_before_unknown_cost(
|
||||
execution_environments_module: Any,
|
||||
) -> None:
|
||||
|
||||
Reference in New Issue
Block a user