Ben FinkleaandClaude Opus 4.5
0b6b9066b3
Bug fixes: - Fix menu button not appearing: Use app.registerExtension() for proper ComfyUI initialization timing. The button now appears in the top bar next to the Manager button. - Fix source_hint incorrectly saved as source_url: Pattern-matched hints like "Lightricks/LTX-Video" are not valid URLs and should not be used for update checks. - Fix deprecated tempfile.mktemp(): Use secure mkstemp() instead to avoid race conditions. - Add os import for file descriptor handling in update_scheduler.py Test improvements: - Update Playwright tests to hover over UI to reveal ComfyUI menu - All 26 tests now passing (17 unit + 9 UI) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
402 lines
14 KiB
Python
402 lines
14 KiB
Python
"""Background scheduler for automatic update checking and downloads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
import time
|
|
import shutil
|
|
import tempfile
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Optional, Callable, Dict, Any, List
|
|
from dataclasses import dataclass, field
|
|
|
|
from .metadata_store import MetadataStore
|
|
from .version_tracker import VersionTracker
|
|
from .hash_utils import compute_sha256_chunked, verify_file_hash
|
|
|
|
|
|
@dataclass
|
|
class DownloadTask:
|
|
"""Represents a pending download task."""
|
|
model_id: int
|
|
source_url: str
|
|
destination: Path
|
|
expected_hash: Optional[str] = None
|
|
status: str = "pending" # pending, downloading, verifying, complete, failed
|
|
progress: float = 0.0
|
|
error: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class SchedulerState:
|
|
"""Current state of the scheduler."""
|
|
running: bool = False
|
|
last_check: Optional[datetime] = None
|
|
next_check: Optional[datetime] = None
|
|
current_downloads: List[DownloadTask] = field(default_factory=list)
|
|
models_checked: int = 0
|
|
updates_found: int = 0
|
|
|
|
|
|
class UpdateScheduler:
|
|
"""Background scheduler for periodic update checks and downloads."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: MetadataStore,
|
|
tracker: VersionTracker,
|
|
check_interval_hours: float = 24.0,
|
|
on_update_found: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
on_download_complete: Optional[Callable[[int, bool], None]] = None,
|
|
):
|
|
"""
|
|
Initialize the update scheduler.
|
|
|
|
Args:
|
|
store: MetadataStore instance
|
|
tracker: VersionTracker instance
|
|
check_interval_hours: Hours between automatic checks
|
|
on_update_found: Callback when update is found (receives model dict)
|
|
on_download_complete: Callback when download finishes (model_id, success)
|
|
"""
|
|
self.store = store
|
|
self.tracker = tracker
|
|
self.check_interval_hours = check_interval_hours
|
|
self.on_update_found = on_update_found
|
|
self.on_download_complete = on_download_complete
|
|
|
|
self._state = SchedulerState()
|
|
self._thread: Optional[threading.Thread] = None
|
|
self._stop_event = threading.Event()
|
|
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
|
|
@property
|
|
def state(self) -> SchedulerState:
|
|
"""Get current scheduler state."""
|
|
return self._state
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
"""Check if scheduler is running."""
|
|
return self._state.running
|
|
|
|
def start(self):
|
|
"""Start the background scheduler thread."""
|
|
if self._state.running:
|
|
return
|
|
|
|
self._stop_event.clear()
|
|
self._thread = threading.Thread(
|
|
target=self._run_loop,
|
|
name="ModelSync-Scheduler",
|
|
daemon=True
|
|
)
|
|
self._state.running = True
|
|
self._thread.start()
|
|
|
|
def stop(self):
|
|
"""Stop the background scheduler."""
|
|
if not self._state.running:
|
|
return
|
|
|
|
self._stop_event.set()
|
|
if self._thread:
|
|
self._thread.join(timeout=5.0)
|
|
self._state.running = False
|
|
|
|
def _run_loop(self):
|
|
"""Main scheduler loop (runs in background thread)."""
|
|
# Create event loop for this thread
|
|
self._loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(self._loop)
|
|
|
|
try:
|
|
while not self._stop_event.is_set():
|
|
# Read current config values (so changes take effect)
|
|
auto_enabled = self.store.get_config("auto_update_enabled", "false")
|
|
interval_hours = float(self.store.get_config("check_interval_hours", "24"))
|
|
|
|
if auto_enabled.lower() == "true":
|
|
# Run the check
|
|
self._loop.run_until_complete(self._check_all_models())
|
|
|
|
# Process any pending auto-updates
|
|
self._loop.run_until_complete(self._download_pending_updates())
|
|
|
|
# Calculate next check time using config value
|
|
self._state.next_check = datetime.now() + timedelta(hours=interval_hours)
|
|
|
|
# Wait for interval or stop signal
|
|
interval_seconds = interval_hours * 3600
|
|
self._stop_event.wait(timeout=interval_seconds)
|
|
|
|
finally:
|
|
self._loop.close()
|
|
self._loop = None
|
|
|
|
async def _check_all_models(self):
|
|
"""Check all models for updates."""
|
|
self._state.last_check = datetime.now()
|
|
self._state.models_checked = 0
|
|
self._state.updates_found = 0
|
|
|
|
# Get all models with source URLs
|
|
models = self.store.get_all_models()
|
|
models_with_urls = [m for m in models if m.get("source_url")]
|
|
|
|
# Batch check for updates
|
|
results = await self.tracker.batch_check(models_with_urls)
|
|
|
|
for model in models_with_urls:
|
|
self._state.models_checked += 1
|
|
model_id = model["id"]
|
|
|
|
if model_id in results:
|
|
remote_info = results[model_id]
|
|
|
|
# Update remote info in database
|
|
updates = {
|
|
"remote_hash": remote_info.get("hash"),
|
|
"remote_size": remote_info.get("size"),
|
|
"remote_modified": remote_info.get("modified"),
|
|
"last_checked": datetime.now().isoformat(),
|
|
}
|
|
|
|
# Check if update is available
|
|
is_update = self.tracker.is_update_available(
|
|
local_hash=model.get("file_hash"),
|
|
local_size=model.get("file_size"),
|
|
remote_hash=remote_info.get("hash"),
|
|
remote_size=remote_info.get("size"),
|
|
)
|
|
|
|
updates["update_available"] = 1 if is_update else 0
|
|
|
|
self.store.update_model(model_id, updates)
|
|
|
|
if is_update:
|
|
self._state.updates_found += 1
|
|
if self.on_update_found:
|
|
# Refresh model data
|
|
updated_model = self.store.get_model(model_id)
|
|
self.on_update_found(updated_model)
|
|
|
|
async def _download_pending_updates(self):
|
|
"""Download updates for models with auto_update enabled."""
|
|
models = self.store.get_auto_update_models()
|
|
models_to_update = [m for m in models if m.get("update_available")]
|
|
|
|
max_concurrent = int(self.store.get_config("max_concurrent_downloads", "2"))
|
|
|
|
for model in models_to_update:
|
|
# Create download task
|
|
task = DownloadTask(
|
|
model_id=model["id"],
|
|
source_url=model["source_url"],
|
|
destination=Path(model["path"]),
|
|
expected_hash=model.get("remote_hash"),
|
|
)
|
|
self._state.current_downloads.append(task)
|
|
|
|
try:
|
|
success = await self._download_model(task)
|
|
if self.on_download_complete:
|
|
self.on_download_complete(model["id"], success)
|
|
finally:
|
|
self._state.current_downloads.remove(task)
|
|
|
|
async def _download_model(self, task: DownloadTask) -> bool:
|
|
"""
|
|
Download a model update.
|
|
|
|
Args:
|
|
task: DownloadTask to process
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
task.status = "downloading"
|
|
|
|
backup_before = self.store.get_config("backup_before_update", "true")
|
|
verify_downloads = self.store.get_config("verify_downloads", "true")
|
|
|
|
backup_path = None
|
|
temp_path = None
|
|
|
|
try:
|
|
# Create backup if enabled
|
|
if backup_before.lower() == "true" and task.destination.exists():
|
|
backup_path = task.destination.with_suffix(
|
|
task.destination.suffix + ".bak"
|
|
)
|
|
shutil.copy2(task.destination, backup_path)
|
|
|
|
# Download to temp file (use mkstemp for security, not deprecated mktemp)
|
|
fd, temp_str = tempfile.mkstemp(suffix=task.destination.suffix)
|
|
os.close(fd) # Close the file descriptor, we'll write to it later
|
|
temp_path = Path(temp_str)
|
|
|
|
# Perform download based on source type
|
|
source_type = self.tracker.detect_source_type(task.source_url)
|
|
|
|
if source_type == "huggingface":
|
|
await self._download_from_huggingface(task, temp_path)
|
|
elif source_type == "civitai":
|
|
await self._download_from_civitai(task, temp_path)
|
|
else:
|
|
task.status = "failed"
|
|
task.error = f"Unsupported source type: {source_type}"
|
|
return False
|
|
|
|
# Verify hash if enabled
|
|
if verify_downloads.lower() == "true" and task.expected_hash:
|
|
task.status = "verifying"
|
|
if not verify_file_hash(temp_path, task.expected_hash):
|
|
task.status = "failed"
|
|
task.error = "Hash verification failed"
|
|
return False
|
|
|
|
# Move to final destination
|
|
shutil.move(str(temp_path), str(task.destination))
|
|
temp_path = None
|
|
|
|
# Update database
|
|
new_hash = compute_sha256_chunked(task.destination)
|
|
new_size = task.destination.stat().st_size
|
|
|
|
self.store.update_model(task.model_id, {
|
|
"file_hash": new_hash,
|
|
"file_size": new_size,
|
|
"update_available": 0,
|
|
"install_date": datetime.now().isoformat(),
|
|
})
|
|
|
|
# Remove backup
|
|
if backup_path and backup_path.exists():
|
|
backup_path.unlink()
|
|
|
|
task.status = "complete"
|
|
task.progress = 1.0
|
|
return True
|
|
|
|
except Exception as e:
|
|
task.status = "failed"
|
|
task.error = str(e)
|
|
|
|
# Restore from backup if available
|
|
if backup_path and backup_path.exists():
|
|
if task.destination.exists():
|
|
task.destination.unlink()
|
|
shutil.move(str(backup_path), str(task.destination))
|
|
|
|
return False
|
|
|
|
finally:
|
|
# Clean up temp file
|
|
if temp_path and temp_path.exists():
|
|
temp_path.unlink()
|
|
|
|
async def _download_from_huggingface(self, task: DownloadTask, dest: Path):
|
|
"""Download file from HuggingFace."""
|
|
try:
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
parsed = self.tracker.parse_huggingface_url(task.source_url)
|
|
if not parsed:
|
|
raise ValueError(f"Could not parse HuggingFace URL: {task.source_url}")
|
|
|
|
repo_id, file_path = parsed
|
|
|
|
# Download to temp location
|
|
downloaded = hf_hub_download(
|
|
repo_id=repo_id,
|
|
filename=file_path,
|
|
local_dir=dest.parent,
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
# Move to expected destination
|
|
shutil.move(downloaded, dest)
|
|
|
|
except ImportError:
|
|
raise ImportError("huggingface_hub not installed")
|
|
|
|
async def _download_from_civitai(self, task: DownloadTask, dest: Path):
|
|
"""Download file from CivitAI."""
|
|
import aiohttp
|
|
|
|
parsed = self.tracker.parse_civitai_url(task.source_url)
|
|
if not parsed:
|
|
raise ValueError(f"Could not parse CivitAI URL: {task.source_url}")
|
|
|
|
# Get download URL
|
|
if "version_id" in parsed:
|
|
download_url = f"https://civitai.com/api/download/models/{parsed['version_id']}"
|
|
else:
|
|
# Need to get version ID first
|
|
async with aiohttp.ClientSession() as session:
|
|
model_url = f"{self.tracker.CIVITAI_API_BASE}/models/{parsed['model_id']}"
|
|
async with session.get(model_url) as resp:
|
|
data = await resp.json()
|
|
version_id = data["modelVersions"][0]["id"]
|
|
download_url = f"https://civitai.com/api/download/models/{version_id}"
|
|
|
|
# Download file
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(download_url) as resp:
|
|
if resp.status != 200:
|
|
raise Exception(f"Download failed with status {resp.status}")
|
|
|
|
total_size = int(resp.headers.get("content-length", 0))
|
|
downloaded = 0
|
|
|
|
with open(dest, "wb") as f:
|
|
async for chunk in resp.content.iter_chunked(8192):
|
|
f.write(chunk)
|
|
downloaded += len(chunk)
|
|
if total_size > 0:
|
|
task.progress = downloaded / total_size
|
|
|
|
async def check_now(self) -> Dict[str, int]:
|
|
"""
|
|
Trigger an immediate check (can be called from main thread).
|
|
|
|
Returns:
|
|
Dict with 'checked' and 'updates_found' counts
|
|
"""
|
|
await self._check_all_models()
|
|
return {
|
|
"checked": self._state.models_checked,
|
|
"updates_found": self._state.updates_found,
|
|
}
|
|
|
|
async def download_update(self, model_id: int) -> bool:
|
|
"""
|
|
Download update for a specific model.
|
|
|
|
Args:
|
|
model_id: Model ID to update
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
model = self.store.get_model(model_id)
|
|
if not model:
|
|
return False
|
|
|
|
task = DownloadTask(
|
|
model_id=model["id"],
|
|
source_url=model["source_url"],
|
|
destination=Path(model["path"]),
|
|
expected_hash=model.get("remote_hash"),
|
|
)
|
|
|
|
self._state.current_downloads.append(task)
|
|
try:
|
|
return await self._download_model(task)
|
|
finally:
|
|
self._state.current_downloads.remove(task)
|