Add conservative remote memory admission

This commit is contained in:
Ken Simpson
2026-08-22 16:22:37 -07:00
parent 8c2911e2f8
commit 19232a9738
7 changed files with 403 additions and 62 deletions
+2
View File
@@ -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
View File
@@ -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[
+1
View File
@@ -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.
+43 -4
View File
@@ -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
View File
@@ -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."""
+99 -1
View File
@@ -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:
+40 -1
View File
@@ -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: