Ben FinkleaandClaude Opus 4.5
9256a1c87c
A custom node extension for ComfyUI that provides: - Model version tracking for HuggingFace and CivitAI - Automatic model updates with configurable scheduling - Multi-GPU synchronization with VRAM-based filtering - Metadata management with hash verification Features: - SQLite database for model metadata storage - REST API endpoints for all operations - Web UI panel integrated into ComfyUI menu - Background scheduler for automatic update checks - rsync-based sync to remote machines - VRAM estimation for model filtering Tests: - 17 unit tests covering all library modules - Playwright integration tests for UI verification Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
457 lines
14 KiB
Python
457 lines
14 KiB
Python
"""Multi-GPU sync manager for model distribution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Optional, List, Dict, Any, Callable, Tuple
|
|
from dataclasses import dataclass, field
|
|
|
|
from .metadata_store import MetadataStore
|
|
|
|
|
|
# VRAM requirement estimates for common model patterns
|
|
VRAM_ESTIMATES: Dict[str, int] = {
|
|
# Large video models
|
|
"wan": 48000,
|
|
"hunyuan": 48000,
|
|
"mochi": 48000,
|
|
"cogvideo": 48000,
|
|
|
|
# Flux variants
|
|
"flux-dev": 24000,
|
|
"flux_dev": 24000,
|
|
"flux-schnell": 12000,
|
|
"flux_schnell": 12000,
|
|
"flux-fill": 24000,
|
|
|
|
# SD variants
|
|
"sdxl": 12000,
|
|
"sd_xl": 12000,
|
|
"sd3": 16000,
|
|
"sd15": 6000,
|
|
"sd1.5": 6000,
|
|
"sd-1": 6000,
|
|
|
|
# Other large models
|
|
"cascade": 20000,
|
|
"pixart": 16000,
|
|
"playground": 16000,
|
|
|
|
# LoRA typically need less
|
|
"lora": 2000,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class SyncResult:
|
|
"""Result of a sync operation."""
|
|
target_name: str
|
|
success: bool
|
|
models_synced: int = 0
|
|
models_skipped: int = 0
|
|
bytes_transferred: int = 0
|
|
duration_seconds: float = 0.0
|
|
error: Optional[str] = None
|
|
details: List[str] = field(default_factory=list)
|
|
|
|
|
|
class SyncManager:
|
|
"""Manages model synchronization to remote machines."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: MetadataStore,
|
|
base_model_path: Optional[str] = None,
|
|
on_progress: Optional[Callable[[str, float], None]] = None,
|
|
):
|
|
"""
|
|
Initialize the sync manager.
|
|
|
|
Args:
|
|
store: MetadataStore instance
|
|
base_model_path: Base path for ComfyUI models directory
|
|
on_progress: Callback(target_name, progress) for progress updates
|
|
"""
|
|
self.store = store
|
|
self.base_model_path = Path(base_model_path) if base_model_path else None
|
|
self.on_progress = on_progress
|
|
|
|
def estimate_vram(self, model: Dict[str, Any]) -> int:
|
|
"""
|
|
Estimate VRAM requirement for a model.
|
|
|
|
Args:
|
|
model: Model dict from database
|
|
|
|
Returns:
|
|
Estimated VRAM in MB
|
|
"""
|
|
filename = model.get("filename", "").lower()
|
|
model_type = model.get("model_type", "").lower()
|
|
|
|
# Check filename patterns
|
|
for pattern, vram in VRAM_ESTIMATES.items():
|
|
if pattern in filename:
|
|
return vram
|
|
|
|
# Check if it's a LoRA by type
|
|
if "lora" in model_type:
|
|
return VRAM_ESTIMATES.get("lora", 2000)
|
|
|
|
# Size-based fallback
|
|
file_size = model.get("file_size", 0)
|
|
if file_size > 0:
|
|
size_gb = file_size / (1024 ** 3)
|
|
# Rough estimate: 2GB VRAM per GB of model file
|
|
return int(size_gb * 2000)
|
|
|
|
# Default: assume it needs moderate VRAM
|
|
return 8000
|
|
|
|
def get_models_for_target(
|
|
self,
|
|
target: Dict[str, Any],
|
|
models: Optional[List[Dict[str, Any]]] = None
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get models filtered by target's VRAM limit.
|
|
|
|
Args:
|
|
target: Sync target dict
|
|
models: Optional list of models (fetches all if not provided)
|
|
|
|
Returns:
|
|
List of models that fit target's VRAM limit
|
|
"""
|
|
if models is None:
|
|
models = self.store.get_all_models()
|
|
|
|
max_vram = target.get("max_vram_mb")
|
|
|
|
filtered = []
|
|
for model in models:
|
|
# Skip if explicitly excluded
|
|
if model.get("exclude_from_sync"):
|
|
continue
|
|
|
|
# Check VRAM requirement
|
|
if max_vram is not None:
|
|
model_vram = model.get("vram_mb") or self.estimate_vram(model)
|
|
if model_vram > max_vram:
|
|
continue
|
|
|
|
filtered.append(model)
|
|
|
|
return filtered
|
|
|
|
def get_relative_path(self, model: Dict[str, Any]) -> str:
|
|
"""
|
|
Get the relative path for a model within the models directory.
|
|
|
|
Args:
|
|
model: Model dict
|
|
|
|
Returns:
|
|
Relative path string
|
|
"""
|
|
model_path = Path(model["path"])
|
|
|
|
if self.base_model_path:
|
|
try:
|
|
return str(model_path.relative_to(self.base_model_path))
|
|
except ValueError:
|
|
pass
|
|
|
|
# Fallback: try to extract from common patterns
|
|
# e.g., /path/to/ComfyUI/models/checkpoints/model.safetensors
|
|
parts = model_path.parts
|
|
for i, part in enumerate(parts):
|
|
if part == "models" and i < len(parts) - 1:
|
|
return str(Path(*parts[i + 1:]))
|
|
|
|
# Last resort: just the filename
|
|
return model_path.name
|
|
|
|
def sync_to_target(
|
|
self,
|
|
target: Dict[str, Any],
|
|
models: Optional[List[Dict[str, Any]]] = None,
|
|
dry_run: bool = False
|
|
) -> SyncResult:
|
|
"""
|
|
Sync models to a target machine using rsync.
|
|
|
|
Args:
|
|
target: Sync target dict
|
|
models: Optional list of models (uses filtered list if not provided)
|
|
dry_run: If True, only simulate the sync
|
|
|
|
Returns:
|
|
SyncResult with details
|
|
"""
|
|
start_time = datetime.now()
|
|
result = SyncResult(target_name=target["name"], success=True)
|
|
|
|
if models is None:
|
|
models = self.get_models_for_target(target)
|
|
|
|
if not models:
|
|
result.details.append("No models to sync")
|
|
return result
|
|
|
|
ssh_host = target.get("ssh_host")
|
|
target_path = target["model_path"]
|
|
|
|
for model in models:
|
|
rel_path = self.get_relative_path(model)
|
|
model_path = model["path"]
|
|
|
|
# Build destination path
|
|
if ssh_host:
|
|
dest = f"{ssh_host}:{target_path}/{rel_path}"
|
|
else:
|
|
dest = f"{target_path}/{rel_path}"
|
|
|
|
# Ensure destination directory exists
|
|
dest_dir = Path(rel_path).parent
|
|
if ssh_host:
|
|
mkdir_cmd = ["ssh", ssh_host, f"mkdir -p '{target_path}/{dest_dir}'"]
|
|
else:
|
|
mkdir_cmd = ["mkdir", "-p", f"{target_path}/{dest_dir}"]
|
|
|
|
try:
|
|
if not dry_run:
|
|
subprocess.run(mkdir_cmd, check=True, capture_output=True)
|
|
except subprocess.CalledProcessError as e:
|
|
result.details.append(f"Failed to create directory for {rel_path}: {e}")
|
|
result.models_skipped += 1
|
|
continue
|
|
|
|
# Build rsync command
|
|
rsync_args = [
|
|
"rsync",
|
|
"-avz", # archive, verbose, compress
|
|
"--progress",
|
|
]
|
|
|
|
if dry_run:
|
|
rsync_args.append("--dry-run")
|
|
|
|
# Add SSH options if remote
|
|
if ssh_host:
|
|
rsync_args.extend(["-e", "ssh"])
|
|
|
|
rsync_args.extend([model_path, dest])
|
|
|
|
try:
|
|
proc = subprocess.run(
|
|
rsync_args,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=3600 # 1 hour timeout per file
|
|
)
|
|
|
|
if proc.returncode == 0:
|
|
result.models_synced += 1
|
|
# Parse transferred bytes from output
|
|
# rsync output: "sent X bytes received Y bytes"
|
|
for line in proc.stdout.split("\n"):
|
|
if "sent" in line and "bytes" in line:
|
|
try:
|
|
sent = int(line.split("sent")[1].split("bytes")[0].strip().replace(",", ""))
|
|
result.bytes_transferred += sent
|
|
except (ValueError, IndexError):
|
|
pass
|
|
else:
|
|
result.details.append(f"rsync failed for {rel_path}: {proc.stderr}")
|
|
result.models_skipped += 1
|
|
|
|
except subprocess.TimeoutExpired:
|
|
result.details.append(f"Timeout syncing {rel_path}")
|
|
result.models_skipped += 1
|
|
except Exception as e:
|
|
result.details.append(f"Error syncing {rel_path}: {e}")
|
|
result.models_skipped += 1
|
|
|
|
# Report progress
|
|
if self.on_progress:
|
|
progress = (result.models_synced + result.models_skipped) / len(models)
|
|
self.on_progress(target["name"], progress)
|
|
|
|
# Update last sync time
|
|
if not dry_run and result.models_synced > 0:
|
|
self.store.update_sync_target(target["id"], {
|
|
"last_sync": datetime.now().isoformat()
|
|
})
|
|
|
|
result.duration_seconds = (datetime.now() - start_time).total_seconds()
|
|
result.success = result.models_skipped == 0
|
|
|
|
return result
|
|
|
|
def sync_to_all_targets(
|
|
self,
|
|
dry_run: bool = False
|
|
) -> List[SyncResult]:
|
|
"""
|
|
Sync models to all enabled targets.
|
|
|
|
Args:
|
|
dry_run: If True, only simulate the sync
|
|
|
|
Returns:
|
|
List of SyncResults for each target
|
|
"""
|
|
targets = self.store.get_all_sync_targets(enabled_only=True)
|
|
results = []
|
|
|
|
for target in targets:
|
|
result = self.sync_to_target(target, dry_run=dry_run)
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
def generate_sync_script(
|
|
self,
|
|
target: Dict[str, Any],
|
|
models: Optional[List[Dict[str, Any]]] = None
|
|
) -> str:
|
|
"""
|
|
Generate a standalone bash script for syncing.
|
|
|
|
Args:
|
|
target: Sync target dict
|
|
models: Optional list of models
|
|
|
|
Returns:
|
|
Bash script content
|
|
"""
|
|
if models is None:
|
|
models = self.get_models_for_target(target)
|
|
|
|
ssh_host = target.get("ssh_host")
|
|
target_path = target["model_path"]
|
|
|
|
lines = [
|
|
"#!/bin/bash",
|
|
"# ModelSync script for target: " + target["name"],
|
|
f"# Generated: {datetime.now().isoformat()}",
|
|
f"# Models: {len(models)}",
|
|
"",
|
|
"set -e",
|
|
"",
|
|
f'TARGET_PATH="{target_path}"',
|
|
]
|
|
|
|
if ssh_host:
|
|
lines.append(f'SSH_HOST="{ssh_host}"')
|
|
lines.append("")
|
|
lines.append("# Ensure remote directories exist")
|
|
else:
|
|
lines.append("")
|
|
lines.append("# Ensure directories exist")
|
|
|
|
# Collect unique directories
|
|
dirs = set()
|
|
for model in models:
|
|
rel_path = self.get_relative_path(model)
|
|
parent = str(Path(rel_path).parent)
|
|
if parent and parent != ".":
|
|
dirs.add(parent)
|
|
|
|
for d in sorted(dirs):
|
|
if ssh_host:
|
|
lines.append(f'ssh "$SSH_HOST" "mkdir -p \'$TARGET_PATH/{d}\'"')
|
|
else:
|
|
lines.append(f'mkdir -p "$TARGET_PATH/{d}"')
|
|
|
|
lines.append("")
|
|
lines.append("# Sync models")
|
|
|
|
for model in models:
|
|
rel_path = self.get_relative_path(model)
|
|
src = model["path"]
|
|
|
|
if ssh_host:
|
|
lines.append(
|
|
f'rsync -avz --progress "{src}" "$SSH_HOST:$TARGET_PATH/{rel_path}"'
|
|
)
|
|
else:
|
|
lines.append(
|
|
f'rsync -avz --progress "{src}" "$TARGET_PATH/{rel_path}"'
|
|
)
|
|
|
|
lines.append("")
|
|
lines.append("echo 'Sync complete!'")
|
|
|
|
return "\n".join(lines)
|
|
|
|
def verify_target_connectivity(self, target: Dict[str, Any]) -> Tuple[bool, str]:
|
|
"""
|
|
Verify SSH connectivity to a target.
|
|
|
|
Args:
|
|
target: Sync target dict
|
|
|
|
Returns:
|
|
Tuple of (success, message)
|
|
"""
|
|
ssh_host = target.get("ssh_host")
|
|
|
|
if not ssh_host:
|
|
# Local target, check if path exists
|
|
target_path = Path(target["model_path"])
|
|
if target_path.exists():
|
|
return True, f"Local path exists: {target_path}"
|
|
else:
|
|
return False, f"Local path does not exist: {target_path}"
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["ssh", "-o", "ConnectTimeout=10", ssh_host, "echo", "OK"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=15
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
return True, f"SSH connection successful to {ssh_host}"
|
|
else:
|
|
return False, f"SSH connection failed: {result.stderr}"
|
|
|
|
except subprocess.TimeoutExpired:
|
|
return False, f"SSH connection timeout to {ssh_host}"
|
|
except Exception as e:
|
|
return False, f"Error connecting to {ssh_host}: {e}"
|
|
|
|
def get_sync_summary(self) -> Dict[str, Any]:
|
|
"""
|
|
Get a summary of sync status across all targets.
|
|
|
|
Returns:
|
|
Dict with summary information
|
|
"""
|
|
targets = self.store.get_all_sync_targets()
|
|
all_models = self.store.get_all_models()
|
|
|
|
summary = {
|
|
"total_models": len(all_models),
|
|
"targets": [],
|
|
}
|
|
|
|
for target in targets:
|
|
eligible_models = self.get_models_for_target(target, all_models)
|
|
|
|
summary["targets"].append({
|
|
"name": target["name"],
|
|
"enabled": bool(target.get("enabled")),
|
|
"max_vram_mb": target.get("max_vram_mb"),
|
|
"eligible_models": len(eligible_models),
|
|
"last_sync": target.get("last_sync"),
|
|
})
|
|
|
|
return summary
|