From 5e85fdb8da0e130dd91645caa40517b6b0d593d7 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 4 Aug 2026 06:01:08 +0200 Subject: [PATCH] perf: optimize LoRA Manager archive hash search and guard SHA-256 fallback - Skip name fallback in search provider runner when request already has sha256 - Add direct SQLite hash lookup in lora_manager_archive with result caching - Add timing and logging for archive search - Refresh HuggingFace author index metadata - Add tests for hash search and provider fallback behavior --- core/services/search_providers.py | 6 +- core/sources/lora_manager_archive.py | 182 ++++++++++++++---------- metadata/huggingface-author-index.json | 4 +- tests/test_lora_manager_archive.py | 22 ++- tests/test_search_service_components.py | 22 +++ 5 files changed, 158 insertions(+), 78 deletions(-) diff --git a/core/services/search_providers.py b/core/services/search_providers.py index b5a9726..4cf4d79 100644 --- a/core/services/search_providers.py +++ b/core/services/search_providers.py @@ -196,7 +196,11 @@ class SearchProviderRunner: self.raise_if_search_cancelled(request, source_key) self.owner.log_search_result(source_key, result) - if not result and request.base_model_context: + if ( + not result + and request.base_model_context + and not getattr(request, "sha256", "") + ): self.raise_if_search_cancelled(request, source_key) self.owner.search_tracker.update( request.progress_id, diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index fd68d9b..17fd459 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -172,50 +172,91 @@ def _load_version_files( files: List[Dict[str, Any]] = [] for row in rows: - try: - data = json.loads(row["data"]) - except Exception: - continue - - if not isinstance(data, dict): - continue - - mirrors = data.get("mirrors") or [] - if not isinstance(mirrors, list): - mirrors = [mirrors] - - available_mirror = next( - ( - mirror - for mirror in mirrors - if isinstance(mirror, dict) and mirror.get("deletedAt") is None - ), - None, - ) - - download_url = data.get("downloadUrl") or ( - available_mirror.get("url") if available_mirror else None - ) - file_name = data.get("name") or ( - available_mirror.get("filename") if available_mirror else None - ) - - files.append( - { - "id": data.get("id"), - "name": file_name, - "type": data.get("type"), - "sizeKB": data.get("sizeKB"), - "downloadUrl": download_url, - "primary": bool(data.get("primary", False)), - "hashes": data.get("hashes", {}), - "metadata": data.get("metadata"), - } - ) - + file_info = _build_file_info(row["data"]) + if file_info: + files.append(file_info) return files +def _build_file_info(raw_data: Any) -> Optional[Dict[str, Any]]: + """Normalize one archived model-file JSON payload.""" + try: + data = json.loads(raw_data) if isinstance(raw_data, str) else raw_data + except Exception: + return None + + if not isinstance(data, dict): + return None + + mirrors = data.get("mirrors") or [] + if not isinstance(mirrors, list): + mirrors = [mirrors] + + available_mirror = next( + ( + mirror + for mirror in mirrors + if isinstance(mirror, dict) and mirror.get("deletedAt") is None + ), + None, + ) + + download_url = data.get("downloadUrl") or ( + available_mirror.get("url") if available_mirror else None + ) + file_name = data.get("name") or ( + available_mirror.get("filename") if available_mirror else None + ) + + return { + "id": data.get("id"), + "name": file_name, + "type": data.get("type"), + "sizeKB": data.get("sizeKB"), + "downloadUrl": download_url, + "primary": bool(data.get("primary", False)), + "hashes": data.get("hashes", {}), + "metadata": data.get("metadata"), + } + + +def _find_hash_match_row( + conn: sqlite3.Connection, requested_sha256: str +) -> Optional[tuple[sqlite3.Row, Dict[str, Any]]]: + """Find an archived file by hash with one filtered SQLite scan.""" + rows = conn.execute( + """ + SELECT + version_id, + id AS file_id, + data AS file_data + FROM model_files + WHERE type = 'Model' + AND data LIKE ? COLLATE NOCASE + ORDER BY id ASC + """, + (f"%{requested_sha256}%",), + ) + + for row in rows: + file_info = _build_file_info(row["file_data"]) + if not file_info: + continue + hashes = file_info.get("hashes") or {} + file_sha256 = normalize_sha256( + file_info.get("sha256") + or hashes.get("SHA256") + or hashes.get("sha256") + ) + if file_sha256 == requested_sha256: + full_row = _load_full_rows_for_versions( + conn, [row["version_id"]] + ).get(row["version_id"]) + if full_row: + return full_row, file_info + return None + + def _query_candidate_rows( conn: sqlite3.Connection, query: str, @@ -686,44 +727,27 @@ def search_lora_manager_archive_for_file( """ requested_sha256 = normalize_sha256(sha256) if requested_sha256: + cache_key = ( + f"hash::{requested_sha256}::{_normalize_model_type(model_type)}::" + f"{_normalize_base_model(base_model_context or '')}" + ) + if cache_key in _search_cache: + return _search_cache[cache_key] + conn = _connect_readonly() if not conn: return None + started_at = time.perf_counter() try: - rows = conn.execute( - """ - SELECT - m.id AS model_id, - m.name AS model_name, - m.type AS model_type, - v.id AS version_id, - v.name AS version_name, - v.base_model AS base_model, - v.position AS position - FROM models m - JOIN model_versions v ON v.model_id = m.id - ORDER BY m.id ASC, COALESCE(v.position, 999999) ASC - """ - ).fetchall() - for row in rows: - files = _load_version_files(conn, row["version_id"]) - matching_file = next( - ( - file_info - for file_info in files - if normalize_sha256( - file_info.get("sha256") - or (file_info.get("hashes") or {}).get("SHA256") - or (file_info.get("hashes") or {}).get("sha256") - ) - == requested_sha256 - ), - None, + match = _find_hash_match_row(conn, requested_sha256) + if match: + row, matching_file = match + result = _build_result_from_row( + conn, row, requested_sha256 ) - if not matching_file: - continue - - result = _build_result_from_row(conn, row) + if not result: + _search_cache[cache_key] = None + return None result.update( { "filename": matching_file.get("name") or result.get("filename", ""), @@ -746,9 +770,19 @@ def search_lora_manager_archive_for_file( 92, confidence=100.0, ) + log.info( + f"LoRA Manager archive hash={requested_sha256} results=1 " + f"elapsed={time.perf_counter() - started_at:.3f}s" + ) + _search_cache[cache_key] = result return result finally: conn.close() + log.info( + f"LoRA Manager archive hash={requested_sha256} results=0 " + f"elapsed={time.perf_counter() - started_at:.3f}s" + ) + _search_cache[cache_key] = None return None search_query = os.path.splitext(get_filename_from_path(filename))[0] diff --git a/metadata/huggingface-author-index.json b/metadata/huggingface-author-index.json index 8dc70a8..0ca1d00 100644 --- a/metadata/huggingface-author-index.json +++ b/metadata/huggingface-author-index.json @@ -3,7 +3,7 @@ "authors": { "Comfy-Org": { "author": "Comfy-Org", - "updated_at": 1785814859.1211712, + "updated_at": 1785815082.3273215, "repo_count": 82, "file_count": 736, "repos": [ @@ -4511,7 +4511,7 @@ }, "Kijai": { "author": "Kijai", - "updated_at": 1785814860.3609285, + "updated_at": 1785815083.420337, "repo_count": 73, "file_count": 932, "repos": [ diff --git a/tests/test_lora_manager_archive.py b/tests/test_lora_manager_archive.py index 5563886..b855ccb 100644 --- a/tests/test_lora_manager_archive.py +++ b/tests/test_lora_manager_archive.py @@ -1,12 +1,15 @@ import unittest import sqlite3 import json +from unittest.mock import patch from core.sources.lora_manager_archive import ( _normalize_model_type, _extract_search_tokens, _query_candidate_rows, _load_full_rows_for_versions, _build_result_from_row, + _find_hash_match_row, + search_lora_manager_archive_for_file, ) class LoraManagerArchiveTests(unittest.TestCase): @@ -103,4 +106,21 @@ class LoraManagerArchiveTests(unittest.TestCase): self.assertEqual(result["sha256"], archive_hash) self.assertEqual(result["hashes"], {"SHA256": archive_hash}) - conn.close() + hash_match = _find_hash_match_row(conn, archive_hash) + self.assertIsNotNone(hash_match) + self.assertEqual(hash_match[0]["version_id"], 10) + self.assertEqual(hash_match[1]["name"], "ghost_v1.safetensors") + + with patch( + "core.sources.lora_manager_archive._connect_readonly", + return_value=conn, + ): + hash_result = search_lora_manager_archive_for_file( + "", + sha256=archive_hash, + ) + + self.assertIsNotNone(hash_result) + self.assertEqual(hash_result["match_type"], "hash") + self.assertEqual(hash_result["confidence"], 100.0) + self.assertEqual(hash_result["sha256"], archive_hash) diff --git a/tests/test_search_service_components.py b/tests/test_search_service_components.py index b36bfca..8e2c3ee 100644 --- a/tests/test_search_service_components.py +++ b/tests/test_search_service_components.py @@ -472,6 +472,28 @@ def test_search_provider_runner_searches_lora_manager_archive(): owner.search_lora_manager_archive_for_file.assert_called_once() +def test_search_provider_runner_skips_name_fallback_for_sha256(): + owner = SimpleNamespace( + logger=MagicMock(), + search_tracker=MagicMock(), + format_log_fields=MagicMock(return_value="file=model.safetensors"), + log_search_result=MagicMock(), + ) + owner.search_tracker.is_cancelled.return_value = False + runner = SearchProviderRunner(owner) + search_fn = MagicMock(return_value=None) + + result = runner.execute_search_with_fallback( + _request(sha256="a" * 64), + "lora_manager_archive", + search_fn, + "LoRA archive", + ) + + assert result is None + search_fn.assert_called_once() + + def test_search_dependencies_require_extension_context(): with pytest.raises(KeyError, match="Missing route dependency: self"): SearchDependencies.from_context(RouteContext({}))