237 lines
9.0 KiB
Python
237 lines
9.0 KiB
Python
"""Asset synchronization helpers for Modal-backed execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Protocol
|
|
|
|
from .settings import ModalSyncSettings, get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SYNC_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".vae"}
|
|
_SKIP_DIRS = {".git", "__pycache__", ".pytest_cache", ".mypy_cache"}
|
|
|
|
|
|
class VolumeBackend(Protocol):
|
|
"""Minimal storage interface needed by the sync engine."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Return whether the remote path already exists."""
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Upload a local file into the remote storage backend."""
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Upload raw bytes into the remote storage backend."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SyncedAsset:
|
|
"""Description of a local asset mirrored into the remote storage root."""
|
|
|
|
local_path: Path
|
|
remote_path: str
|
|
sha256: str
|
|
|
|
|
|
class LocalMirrorVolume:
|
|
"""Simple filesystem-backed volume used for tests and dry runs."""
|
|
|
|
def __init__(self, root: Path) -> None:
|
|
"""Initialize the local mirror volume root."""
|
|
self.root = root
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Return whether a file already exists in the local mirror."""
|
|
return self._resolve(remote_path).exists()
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Copy a local file into the mirror volume."""
|
|
target = self._resolve(remote_path)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
target.write_bytes(local_path.read_bytes())
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Write bytes into the mirror volume."""
|
|
target = self._resolve(remote_path)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
target.write_bytes(payload)
|
|
|
|
def _resolve(self, remote_path: str) -> Path:
|
|
"""Resolve a remote storage path relative to the mirror root."""
|
|
return self.root / remote_path.lstrip("/")
|
|
|
|
|
|
@dataclass
|
|
class ModalAssetSyncEngine:
|
|
"""Content-addressable storage sync engine for files and custom nodes."""
|
|
|
|
volume: VolumeBackend
|
|
settings: ModalSyncSettings
|
|
|
|
@classmethod
|
|
def from_environment(cls, settings: ModalSyncSettings | None = None) -> "ModalAssetSyncEngine":
|
|
"""Create a sync engine using the local mirror backend by default."""
|
|
resolved_settings = settings or get_settings()
|
|
backend = LocalMirrorVolume(resolved_settings.local_storage_root)
|
|
return cls(volume=backend, settings=resolved_settings)
|
|
|
|
def sync_file(self, local_path: Path, remote_folder: str = "/assets") -> SyncedAsset:
|
|
"""Sync a file into content-addressable remote storage."""
|
|
resolved_path = local_path.expanduser().resolve()
|
|
if not resolved_path.is_file():
|
|
raise FileNotFoundError(f"Asset not found: {resolved_path}")
|
|
|
|
sha256 = self._hash_file(resolved_path)
|
|
marker_path = f"/hashes/{sha256}.done"
|
|
remote_path = f"{remote_folder.rstrip('/')}/{sha256}_{resolved_path.name}"
|
|
|
|
if not self.volume.exists(marker_path):
|
|
logger.info("Syncing %s to %s", resolved_path, remote_path)
|
|
self.volume.put_file(resolved_path, remote_path)
|
|
self.volume.put_bytes(
|
|
json.dumps({"source": str(resolved_path), "remote_path": remote_path}).encode(
|
|
"utf-8"
|
|
),
|
|
marker_path,
|
|
)
|
|
else:
|
|
logger.debug("Asset already mirrored for sha=%s", sha256)
|
|
|
|
return SyncedAsset(local_path=resolved_path, remote_path=remote_path, sha256=sha256)
|
|
|
|
def sync_prompt_inputs(self, inputs: dict[str, Any]) -> tuple[dict[str, Any], list[SyncedAsset]]:
|
|
"""Rewrite file-like prompt inputs to mirrored storage paths."""
|
|
synced_assets: list[SyncedAsset] = []
|
|
|
|
def rewrite(value: Any) -> Any:
|
|
if isinstance(value, str):
|
|
maybe_path = self._resolve_model_path(value)
|
|
if maybe_path is not None:
|
|
synced_asset = self.sync_file(maybe_path)
|
|
synced_assets.append(synced_asset)
|
|
return synced_asset.remote_path
|
|
return value
|
|
if isinstance(value, list):
|
|
return [rewrite(item) for item in value]
|
|
if isinstance(value, dict):
|
|
return {str(key): rewrite(item) for key, item in value.items()}
|
|
return value
|
|
|
|
return rewrite(inputs), synced_assets
|
|
|
|
def sync_custom_nodes_directory(self) -> SyncedAsset | None:
|
|
"""Zip and mirror the local custom_nodes directory when available."""
|
|
custom_nodes_dir = self.settings.custom_nodes_dir
|
|
if custom_nodes_dir is None or not custom_nodes_dir.exists():
|
|
logger.info("No custom_nodes directory detected for mirroring.")
|
|
return None
|
|
|
|
directory_hash = self._hash_directory(custom_nodes_dir)
|
|
archive_name = self.settings.custom_nodes_archive_name
|
|
marker_path = f"/hashes/custom_nodes_{directory_hash}.done"
|
|
remote_path = f"/custom_nodes/{directory_hash}_{archive_name}"
|
|
|
|
if self.volume.exists(marker_path):
|
|
return SyncedAsset(
|
|
local_path=custom_nodes_dir,
|
|
remote_path=remote_path,
|
|
sha256=directory_hash,
|
|
)
|
|
|
|
archive_path = self._create_archive(custom_nodes_dir)
|
|
try:
|
|
logger.info("Syncing custom_nodes bundle from %s to %s", custom_nodes_dir, remote_path)
|
|
self.volume.put_file(archive_path, remote_path)
|
|
self.volume.put_bytes(
|
|
json.dumps({"source": str(custom_nodes_dir), "remote_path": remote_path}).encode(
|
|
"utf-8"
|
|
),
|
|
marker_path,
|
|
)
|
|
finally:
|
|
archive_path.unlink(missing_ok=True)
|
|
|
|
return SyncedAsset(local_path=custom_nodes_dir, remote_path=remote_path, sha256=directory_hash)
|
|
|
|
def _resolve_model_path(self, value: str) -> Path | None:
|
|
"""Resolve a prompt string into a local model file path when possible."""
|
|
path = Path(value).expanduser()
|
|
if path.suffix.lower() not in _SYNC_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 self.settings.comfyui_root is not None:
|
|
candidate = self.settings.comfyui_root / value
|
|
if candidate.is_file():
|
|
return candidate.resolve()
|
|
|
|
return None
|
|
|
|
def _hash_file(self, path: Path) -> str:
|
|
"""Compute the SHA256 digest for a file."""
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
def _hash_directory(self, path: Path) -> str:
|
|
"""Compute a stable SHA256 digest for a directory tree."""
|
|
digest = hashlib.sha256()
|
|
for child in sorted(self._iter_files(path), key=lambda item: item.relative_to(path).as_posix()):
|
|
relative_path = child.relative_to(path).as_posix()
|
|
digest.update(relative_path.encode("utf-8"))
|
|
digest.update(b"\0")
|
|
digest.update(self._hash_file(child).encode("ascii"))
|
|
digest.update(b"\0")
|
|
return digest.hexdigest()
|
|
|
|
def _create_archive(self, path: Path) -> Path:
|
|
"""Create a zip archive for the given directory tree."""
|
|
with tempfile.NamedTemporaryFile(
|
|
prefix="comfy-modal-custom-nodes-",
|
|
suffix=".zip",
|
|
delete=False,
|
|
) as handle:
|
|
archive_path = Path(handle.name)
|
|
|
|
with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
|
|
for child in sorted(self._iter_files(path), key=lambda item: item.relative_to(path).as_posix()):
|
|
archive.write(child, arcname=child.relative_to(path))
|
|
|
|
return archive_path
|
|
|
|
def _iter_files(self, path: Path) -> list[Path]:
|
|
"""Yield files from a directory tree while skipping cache folders."""
|
|
files: list[Path] = []
|
|
for root, dirnames, filenames in os.walk(path):
|
|
dirnames[:] = [name for name in dirnames if name not in _SKIP_DIRS]
|
|
for filename in filenames:
|
|
files.append(Path(root) / filename)
|
|
return files
|