refactor: centralize hash calculation callbacks and catalog backup logic
This commit is contained in:
+47
-45
@@ -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
@@ -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):
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user