296 lines
11 KiB
Python
296 lines
11 KiB
Python
"""Runtime configuration helpers for ComfyUI Modal-Sync."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SETTINGS_ENV_KEYS = (
|
|
"COMFYUI_ROOT",
|
|
"COMFY_MODAL_COMFYUI_ROOT",
|
|
"COMFY_MODAL_CUSTOM_NODES_DIR",
|
|
"COMFY_MODAL_EXECUTION_MODE",
|
|
"COMFY_MODAL_SYNC_CUSTOM_NODES",
|
|
"COMFY_MODAL_LOCAL_STORAGE_ROOT",
|
|
"COMFY_MODAL_APP_NAME",
|
|
"COMFY_MODAL_AUTO_DEPLOY",
|
|
"COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK",
|
|
"COMFY_MODAL_ENABLE_MEMORY_SNAPSHOT",
|
|
"COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT",
|
|
"COMFY_MODAL_VOLUME_NAME",
|
|
"COMFY_MODAL_ROUTE_PATH",
|
|
"COMFY_MODAL_MARKER_PROPERTY",
|
|
"COMFY_MODAL_REMOTE_STORAGE_ROOT",
|
|
"COMFY_MODAL_CUSTOM_NODES_ARCHIVE",
|
|
"COMFY_MODAL_INTERRUPT_DICT_NAME",
|
|
"COMFY_MODAL_NODE_CACHE_DICT_NAME",
|
|
"COMFY_MODAL_SESSION_BRIDGE_DICT_NAME",
|
|
"COMFY_MODAL_SYNC_INDEX_DICT_NAME",
|
|
"COMFY_MODAL_SNAPSHOT_PROFILE_DICT_NAME",
|
|
"COMFY_MODAL_NODE_CACHE_MAX_BYTES",
|
|
"COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR",
|
|
"COMFY_MODAL_GPU",
|
|
"COMFY_MODAL_MIN_CONTAINERS",
|
|
"COMFY_MODAL_MAX_CONTAINERS",
|
|
"COMFY_MODAL_BUFFER_CONTAINERS",
|
|
"COMFY_MODAL_SCALEDOWN_WINDOW",
|
|
"COMFY_MODAL_ENABLE_LOADER_PREWARM",
|
|
"COMFY_MODAL_MAX_LOADER_PREWARMS_PER_COMPONENT",
|
|
"COMFY_MODAL_ENABLE_PROACTIVE_WARMUP",
|
|
"COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS",
|
|
"COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS",
|
|
"COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS",
|
|
"COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModalSyncSettings:
|
|
"""Resolved configuration for local and remote Modal-Sync operations."""
|
|
|
|
app_name: str
|
|
auto_deploy: bool
|
|
allow_ephemeral_fallback: bool
|
|
enable_memory_snapshot: bool
|
|
enable_gpu_memory_snapshot: bool
|
|
execution_mode: str
|
|
sync_custom_nodes: bool
|
|
volume_name: str
|
|
route_path: str
|
|
marker_property: str
|
|
local_storage_root: Path
|
|
remote_storage_root: str
|
|
custom_nodes_archive_name: str
|
|
comfyui_root: Path | None
|
|
custom_nodes_dir: Path | None
|
|
interrupt_dict_name: str = "comfy-modal-sync-interrupts"
|
|
node_output_cache_dict_name: str = "comfy-modal-sync-node-cache"
|
|
session_bridge_dict_name: str = "comfy-modal-sync-session-bridges"
|
|
sync_index_dict_name: str = "comfy-modal-sync-sync-index"
|
|
snapshot_profile_dict_name: str = "comfy-modal-sync-snapshot-profiles"
|
|
node_output_cache_max_bytes: int = 5 * 1024 * 1024
|
|
terminate_container_on_error: bool = True
|
|
modal_gpu: str = "A100"
|
|
scaledown_window_seconds: int = 600
|
|
min_containers: int = 0
|
|
max_containers: int | None = None
|
|
buffer_containers: int | None = None
|
|
enable_proactive_warmup: bool = True
|
|
enable_loader_prewarm: bool = True
|
|
proactive_warmup_head_start_seconds: float = 2.0
|
|
remote_cancel_grace_seconds: float = 2.0
|
|
remote_cancel_restart_seconds: float = 1.0
|
|
stream_remote_container_logs: bool = False
|
|
|
|
|
|
def _read_path_env(name: str) -> Path | None:
|
|
"""Resolve an environment variable into a path when present."""
|
|
value = os.getenv(name)
|
|
if not value:
|
|
return None
|
|
return Path(value).expanduser().resolve()
|
|
|
|
|
|
def _read_bool_env(name: str) -> bool | None:
|
|
"""Resolve an environment variable into a boolean when present."""
|
|
value = os.getenv(name)
|
|
if value is None:
|
|
return None
|
|
|
|
normalized = value.strip().lower()
|
|
if normalized in {"1", "true", "yes", "on"}:
|
|
return True
|
|
if normalized in {"0", "false", "no", "off"}:
|
|
return False
|
|
raise ValueError(f"Environment variable {name} must be a boolean, got {value!r}.")
|
|
|
|
|
|
def _discover_repo_root() -> Path:
|
|
"""Return the repository root containing this module."""
|
|
return Path(__file__).resolve().parent
|
|
|
|
|
|
def _read_int_env(name: str, default: int) -> int:
|
|
"""Resolve an environment variable into an integer when present."""
|
|
value = os.getenv(name)
|
|
if value is None:
|
|
return default
|
|
try:
|
|
return int(value)
|
|
except ValueError as exc:
|
|
raise ValueError(f"Environment variable {name} must be an integer, got {value!r}.") from exc
|
|
|
|
|
|
def _read_optional_int_env(name: str) -> int | None:
|
|
"""Resolve an environment variable into an optional integer when present."""
|
|
value = os.getenv(name)
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return int(value)
|
|
except ValueError as exc:
|
|
raise ValueError(f"Environment variable {name} must be an integer, got {value!r}.") from exc
|
|
|
|
|
|
def _read_float_env(name: str, default: float) -> float:
|
|
"""Resolve an environment variable into a float when present."""
|
|
value = os.getenv(name)
|
|
if value is None:
|
|
return default
|
|
try:
|
|
return float(value)
|
|
except ValueError as exc:
|
|
raise ValueError(f"Environment variable {name} must be a float, got {value!r}.") from exc
|
|
|
|
|
|
def _looks_like_comfyui_root(candidate: Path) -> bool:
|
|
"""Return whether a path appears to be a ComfyUI checkout root."""
|
|
return (candidate / "main.py").exists() and (candidate / "nodes.py").exists()
|
|
|
|
|
|
def _discover_comfyui_root(repo_root: Path) -> Path | None:
|
|
"""Locate the local ComfyUI checkout used for tests and path resolution."""
|
|
env_root = _read_path_env("COMFYUI_ROOT")
|
|
if env_root is not None:
|
|
return env_root
|
|
|
|
modal_env_root = _read_path_env("COMFY_MODAL_COMFYUI_ROOT")
|
|
if modal_env_root is not None:
|
|
return modal_env_root
|
|
|
|
if repo_root.parent.name == "custom_nodes":
|
|
install_root = repo_root.parent.parent.resolve()
|
|
if _looks_like_comfyui_root(install_root):
|
|
return install_root
|
|
|
|
default_root = Path.home() / "git" / "ComfyUI"
|
|
if _looks_like_comfyui_root(default_root):
|
|
return default_root.resolve()
|
|
|
|
return None
|
|
|
|
|
|
def _discover_custom_nodes_dir(repo_root: Path, comfyui_root: Path | None) -> Path | None:
|
|
"""Locate the custom_nodes directory that should be mirrored to Modal."""
|
|
env_dir = _read_path_env("COMFY_MODAL_CUSTOM_NODES_DIR")
|
|
if env_dir is not None:
|
|
return env_dir
|
|
|
|
if repo_root.parent.name == "custom_nodes":
|
|
return repo_root.parent.resolve()
|
|
|
|
if comfyui_root is not None:
|
|
candidate = comfyui_root / "custom_nodes"
|
|
if candidate.exists():
|
|
return candidate.resolve()
|
|
|
|
return None
|
|
|
|
|
|
def _settings_env_signature() -> tuple[tuple[str, str | None], ...]:
|
|
"""Return the environment values that affect resolved Modal-Sync settings."""
|
|
return tuple((key, os.getenv(key)) for key in _SETTINGS_ENV_KEYS)
|
|
|
|
|
|
@lru_cache(maxsize=8)
|
|
def _get_settings_cached(
|
|
env_signature: tuple[tuple[str, str | None], ...],
|
|
) -> ModalSyncSettings:
|
|
"""Return cached extension settings for one environment signature."""
|
|
del env_signature
|
|
repo_root = _discover_repo_root()
|
|
comfyui_root = _discover_comfyui_root(repo_root)
|
|
custom_nodes_dir = _discover_custom_nodes_dir(repo_root, comfyui_root)
|
|
execution_mode = os.getenv("COMFY_MODAL_EXECUTION_MODE", "local").strip().lower()
|
|
sync_custom_nodes = _read_bool_env("COMFY_MODAL_SYNC_CUSTOM_NODES")
|
|
if sync_custom_nodes is None:
|
|
sync_custom_nodes = execution_mode != "local"
|
|
local_storage_root = (
|
|
_read_path_env("COMFY_MODAL_LOCAL_STORAGE_ROOT")
|
|
or Path("/tmp/comfyui-modal-sync-storage")
|
|
)
|
|
|
|
settings = ModalSyncSettings(
|
|
app_name=(app_name := os.getenv("COMFY_MODAL_APP_NAME", "comfy-modal-sync")),
|
|
auto_deploy=_read_bool_env("COMFY_MODAL_AUTO_DEPLOY") is not False,
|
|
allow_ephemeral_fallback=_read_bool_env("COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK") or False,
|
|
enable_memory_snapshot=_read_bool_env("COMFY_MODAL_ENABLE_MEMORY_SNAPSHOT") is not False,
|
|
enable_gpu_memory_snapshot=_read_bool_env("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT")
|
|
is not False,
|
|
execution_mode=execution_mode,
|
|
sync_custom_nodes=sync_custom_nodes,
|
|
volume_name=os.getenv("COMFY_MODAL_VOLUME_NAME", "comfy-universal-storage"),
|
|
route_path=os.getenv("COMFY_MODAL_ROUTE_PATH", "/modal/queue_prompt"),
|
|
marker_property=os.getenv("COMFY_MODAL_MARKER_PROPERTY", "is_modal_remote"),
|
|
local_storage_root=local_storage_root.resolve(),
|
|
remote_storage_root=os.getenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", "/storage"),
|
|
custom_nodes_archive_name=os.getenv(
|
|
"COMFY_MODAL_CUSTOM_NODES_ARCHIVE",
|
|
"custom_nodes_bundle.zip",
|
|
),
|
|
comfyui_root=comfyui_root,
|
|
custom_nodes_dir=custom_nodes_dir,
|
|
interrupt_dict_name=os.getenv(
|
|
"COMFY_MODAL_INTERRUPT_DICT_NAME",
|
|
f"{app_name}-interrupts",
|
|
),
|
|
node_output_cache_dict_name=os.getenv(
|
|
"COMFY_MODAL_NODE_CACHE_DICT_NAME",
|
|
f"{app_name}-node-cache",
|
|
),
|
|
session_bridge_dict_name=os.getenv(
|
|
"COMFY_MODAL_SESSION_BRIDGE_DICT_NAME",
|
|
f"{app_name}-session-bridges",
|
|
),
|
|
snapshot_profile_dict_name=os.getenv(
|
|
"COMFY_MODAL_SNAPSHOT_PROFILE_DICT_NAME",
|
|
f"{app_name}-snapshot-profiles",
|
|
),
|
|
sync_index_dict_name=os.getenv(
|
|
"COMFY_MODAL_SYNC_INDEX_DICT_NAME",
|
|
f"{app_name}-sync-index",
|
|
),
|
|
node_output_cache_max_bytes=_read_int_env(
|
|
"COMFY_MODAL_NODE_CACHE_MAX_BYTES",
|
|
5 * 1024 * 1024,
|
|
),
|
|
terminate_container_on_error=_read_bool_env("COMFY_MODAL_TERMINATE_CONTAINER_ON_ERROR")
|
|
is not False,
|
|
modal_gpu=os.getenv("COMFY_MODAL_GPU", "A100").strip() or "A100",
|
|
scaledown_window_seconds=_read_int_env("COMFY_MODAL_SCALEDOWN_WINDOW", 600),
|
|
min_containers=_read_int_env("COMFY_MODAL_MIN_CONTAINERS", 0),
|
|
max_containers=_read_optional_int_env("COMFY_MODAL_MAX_CONTAINERS"),
|
|
buffer_containers=_read_optional_int_env("COMFY_MODAL_BUFFER_CONTAINERS"),
|
|
enable_proactive_warmup=_read_bool_env("COMFY_MODAL_ENABLE_PROACTIVE_WARMUP") is not False,
|
|
enable_loader_prewarm=_read_bool_env("COMFY_MODAL_ENABLE_LOADER_PREWARM") is not False,
|
|
proactive_warmup_head_start_seconds=_read_float_env(
|
|
"COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS",
|
|
2.0,
|
|
),
|
|
remote_cancel_grace_seconds=_read_float_env(
|
|
"COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS",
|
|
2.0,
|
|
),
|
|
remote_cancel_restart_seconds=_read_float_env(
|
|
"COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS",
|
|
1.0,
|
|
),
|
|
stream_remote_container_logs=_read_bool_env("COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS")
|
|
or False,
|
|
)
|
|
logger.debug("Resolved Modal-Sync settings: %s", settings)
|
|
return settings
|
|
|
|
|
|
def get_settings() -> ModalSyncSettings:
|
|
"""Return extension settings derived from the current environment."""
|
|
return _get_settings_cached(_settings_env_signature())
|
|
|
|
|
|
get_settings.cache_clear = _get_settings_cached.cache_clear # type: ignore[attr-defined]
|