refactor: centralize hash calculation callbacks and catalog backup logic

This commit is contained in:
Dariusz L
2026-07-04 00:18:41 +02:00
parent 5de1500537
commit cb666902e5
4 changed files with 93 additions and 69 deletions
+47 -45
View File
@@ -724,8 +724,7 @@ class ModelResolverExtension:
# cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed
class HashCalculationCancelled(Exception):
pass
from .core.path_utils import HashCalculationCancelled
def resolve_hash_file_request(data):
import os as _os
@@ -799,14 +798,10 @@ class ModelResolverExtension:
return resolved_metadata_path, metadata_updated
def calculate_sha256_with_progress(normalized_path, progress_id=""):
import hashlib
import os as _os
from .core.path_utils import calculate_file_sha256 as _calculate_file_sha256_core
total_bytes = max(0, _os.path.getsize(normalized_path))
bytes_read = 0
sha256_hash = hashlib.sha256()
last_update = 0.0
chunk_size = 1024 * 1024 * 4
self.hash_tracker.update(
progress_id,
@@ -818,45 +813,52 @@ class ModelResolverExtension:
total_bytes=total_bytes,
)
with open(normalized_path, "rb") as handle:
for chunk in iter(lambda: handle.read(chunk_size), b""):
if not chunk:
continue
sha256_hash.update(chunk)
bytes_read += len(chunk)
if self.hash_tracker.is_cancelled(progress_id):
percent = 0 if total_bytes <= 0 else min(
98,
(bytes_read / total_bytes) * 98,
)
self.hash_tracker.update(
progress_id,
status="cancelled",
stage="cancelled",
message="Hash calculation cancelled",
percent=percent,
bytes_read=bytes_read,
total_bytes=total_bytes,
)
raise HashCalculationCancelled()
now = time.time()
if now - last_update >= 0.15 or bytes_read >= total_bytes:
percent = 98 if total_bytes <= 0 else min(
98,
(bytes_read / total_bytes) * 98,
)
self.hash_tracker.update(
progress_id,
status="running",
stage="hashing",
message="Calculating SHA256...",
percent=percent,
bytes_read=bytes_read,
total_bytes=total_bytes,
)
last_update = now
_bytes_read = [0]
_last_update = [0.0]
return sha256_hash.hexdigest()
def on_progress(bytes_read, total_bytes):
_bytes_read[0] = bytes_read
now = time.time()
if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes:
percent = 98 if total_bytes <= 0 else min(
98,
(bytes_read / total_bytes) * 98,
)
self.hash_tracker.update(
progress_id,
status="running",
stage="hashing",
message="Calculating SHA256...",
percent=percent,
bytes_read=bytes_read,
total_bytes=total_bytes,
)
_last_update[0] = now
def is_cancelled():
if self.hash_tracker.is_cancelled(progress_id):
percent = 0 if total_bytes <= 0 else min(
98,
(_bytes_read[0] / total_bytes) * 98,
)
self.hash_tracker.update(
progress_id,
status="cancelled",
stage="cancelled",
message="Hash calculation cancelled",
percent=percent,
bytes_read=_bytes_read[0],
total_bytes=total_bytes,
)
return True
return False
return _calculate_file_sha256_core(
normalized_path,
chunk_size=1024 * 1024 * 4,
on_progress=on_progress,
is_cancelled=is_cancelled,
)
@routes.post("/model_resolver/calculate-file-hash")
@json_api_endpoint("calculate-file-hash")
+43 -3
View File
@@ -8,7 +8,7 @@ and directory matching to be shared across all modules.
import os
import json
import hashlib
from typing import Any, Tuple, Optional
from typing import Any, Tuple, Optional, Callable
_log = None
@@ -178,22 +178,62 @@ METADATA_DIR = os.path.join(
)
def calculate_file_sha256(file_path: str, chunk_size: int = 131072) -> Optional[str]:
"""Calculate SHA256 hex digest for an existing local file safely."""
class HashCalculationCancelled(Exception):
"""Exception raised when file SHA256 calculation is cancelled by the caller."""
pass
def calculate_file_sha256(
file_path: str,
chunk_size: int = 131072,
on_progress: Optional[Callable[[int, int], None]] = None,
is_cancelled: Optional[Callable[[], bool]] = None,
) -> Optional[str]:
"""Calculate SHA256 hex digest for an existing local file safely, supporting progress callbacks and cancellation check."""
if not file_path or not os.path.exists(file_path):
return None
sha256_hash = hashlib.sha256()
try:
total_bytes = os.path.getsize(file_path) if on_progress else 0
bytes_read = 0
with open(file_path, "rb") as f:
for byte_block in iter(lambda: f.read(chunk_size), b""):
if is_cancelled and is_cancelled():
raise HashCalculationCancelled("Hash calculation cancelled")
if byte_block:
sha256_hash.update(byte_block)
bytes_read += len(byte_block)
if on_progress:
on_progress(bytes_read, total_bytes)
return sha256_hash.hexdigest()
except HashCalculationCancelled:
raise
except Exception as e:
_get_log().error(f"Error computing hash for {file_path}: {e}")
return None
def save_catalog_with_backup(
data_path: str, data: Any, meta_path: str, meta: Any, indent: int = 2
) -> None:
"""Backup an existing catalog file to .bak, then atomically write the new data and its metadata."""
import shutil
# Ensure directory exists
dir_name = os.path.dirname(data_path)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
# Backup existing data file
if os.path.exists(data_path):
shutil.copy2(data_path, f"{data_path}.bak")
# Write data file atomically
write_json_atomic(data_path, data, indent=indent)
# Write metadata file atomically
write_json_atomic(meta_path, meta, indent=indent)
def read_json_safe(file_path: str, default: Any = None) -> Any:
"""Safely read and parse a JSON file, returning a default value on error."""
if not file_path or not os.path.exists(file_path):
+2 -6
View File
@@ -19,7 +19,7 @@ from ..log_system import create_module_logger
log = create_module_logger(__name__)
from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic
from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic, save_catalog_with_backup
MODEL_LIST_FILE = os.path.join(METADATA_DIR, "model-list.json")
MODEL_LIST_META_FILE = os.path.join(METADATA_DIR, "model-list.meta.json")
MODEL_LIST_SOURCE_URL = (
@@ -152,10 +152,6 @@ def update_model_list_from_remote() -> Dict[str, Any]:
if not isinstance(models, list) or not models:
raise ValueError("Downloaded model-list.json does not contain a non-empty models list")
if os.path.exists(MODEL_LIST_FILE):
shutil.copy2(MODEL_LIST_FILE, f"{MODEL_LIST_FILE}.bak")
write_json_atomic(MODEL_LIST_FILE, data, indent=2)
meta = {
"source": "Comfy-Org/ComfyUI-Manager",
"source_url": MODEL_LIST_SOURCE_URL,
@@ -166,7 +162,7 @@ def update_model_list_from_remote() -> Dict[str, Any]:
"model_count": len(models),
"updated_at": _utc_now_iso(),
}
write_json_atomic(MODEL_LIST_META_FILE, meta, indent=2)
save_catalog_with_backup(MODEL_LIST_FILE, data, MODEL_LIST_META_FILE, meta, indent=2)
reload_model_list()
return {
+1 -15
View File
@@ -349,17 +349,6 @@ def update_base_models_from_remote() -> Dict[str, Any]:
new_added_count += 1
new_added_names.append(name)
# Write updated models to base-models.json atomically
os.makedirs(METADATA_DIR, exist_ok=True)
# Backup existing
if os.path.exists(BASE_MODELS_FILE):
shutil.copy2(BASE_MODELS_FILE, f"{BASE_MODELS_FILE}.bak")
# Write atomically
write_json_atomic(BASE_MODELS_FILE, {"base_models": updated_models}, indent=2)
# Write meta file
now_str = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
meta = {
"updated_at": now_str,
@@ -368,10 +357,7 @@ def update_base_models_from_remote() -> Dict[str, Any]:
"new_models_added": new_added_count,
"new_models_added_list": new_added_names
}
write_json_atomic(BASE_MODELS_META_FILE, meta, indent=2)
# Force reload database in-memory cache
save_catalog_with_backup(BASE_MODELS_FILE, {"base_models": updated_models}, BASE_MODELS_META_FILE, meta, indent=2)
reload_databases()
return {