From ee233548e5b51a40da67097183551064d7986a9b Mon Sep 17 00:00:00 2001 From: Will Miao Date: Thu, 27 Aug 2026 07:58:48 +0800 Subject: [PATCH] fix(recipes): enforce batch-import concurrency bound and harden ingest errors (#1085) Address the rate-limit flood and secondary errors seen during large recipe ingestion (example-images directory import): - batch import: share one adaptive-concurrency semaphore across the whole batch (previously each item got a fresh semaphore, so the min/max concurrency bounds never applied and every item ran concurrently); synchronize the shared semaphore capacity after each completed item. - comfy parser: guard ckpt_name against list/None values so re.search no longer raises TypeError and fails the whole image import. - civarchive client: normalize empty-string failure payloads to "Request failed" and treat a missing payload as an error, fixing the "'NoneType' object has no attribute 'get'" crash. - civarchive client: log connectivity-guard offline-cooldown short-circuits at DEBUG instead of one ERROR per request. --- py/recipes/parsers/comfy.py | 46 ++++++++----- py/services/batch_import_service.py | 40 +++++++++++- py/services/civarchive_client.py | 50 ++++++++++++-- tests/services/test_batch_import_service.py | 59 +++++++++++++++++ tests/services/test_civarchive_client.py | 47 +++++++++++++ tests/services/test_comfy_metadata_parser.py | 69 ++++++++++++++++++++ 6 files changed, 285 insertions(+), 26 deletions(-) diff --git a/py/recipes/parsers/comfy.py b/py/recipes/parsers/comfy.py index 678ad075..dced6fdd 100644 --- a/py/recipes/parsers/comfy.py +++ b/py/recipes/parsers/comfy.py @@ -40,24 +40,34 @@ class ComfyMetadataParser(RecipeMetadataParser): checkpoint_node = next(iter(checkpoint_nodes.values())) if 'inputs' in checkpoint_node and 'ckpt_name' in checkpoint_node['inputs']: checkpoint_name = checkpoint_node['inputs']['ckpt_name'] - checkpoint_match = re.search(r'civitai:(\d+)@(\d+)', checkpoint_name) - if checkpoint_match: - checkpoint_id = checkpoint_match.group(1) - checkpoint_version_id = checkpoint_match.group(2) - checkpoint = { - 'id': checkpoint_version_id, - 'modelId': checkpoint_id, - 'name': f"Checkpoint {checkpoint_id}", - 'version': '', - 'type': 'checkpoint' - } - if metadata_provider: - try: - civitai_info_tuple = await metadata_provider.get_model_version_info(checkpoint_version_id) - civitai_info, _ = civitai_info_tuple if isinstance(civitai_info_tuple, tuple) else (civitai_info_tuple, None) - checkpoint = await self.populate_checkpoint_from_civitai(checkpoint, civitai_info) - except Exception as e: - logger.error(f"Error fetching Civitai info for checkpoint: {e}") + # Some ComfyUI workflows serialize ckpt_name as a + # single-element list (e.g. ["model.safetensors"]) or leave + # the value unset (None). Neither is a string, so skip the + # CivitAI-URN lookup instead of crashing re.search with a + # TypeError that fails the whole image import. + if isinstance(checkpoint_name, list): + checkpoint_name = ( + checkpoint_name[0] if checkpoint_name else None + ) + if isinstance(checkpoint_name, str): + checkpoint_match = re.search(r'civitai:(\d+)@(\d+)', checkpoint_name) + if checkpoint_match: + checkpoint_id = checkpoint_match.group(1) + checkpoint_version_id = checkpoint_match.group(2) + checkpoint = { + 'id': checkpoint_version_id, + 'modelId': checkpoint_id, + 'name': f"Checkpoint {checkpoint_id}", + 'version': '', + 'type': 'checkpoint' + } + if metadata_provider: + try: + civitai_info_tuple = await metadata_provider.get_model_version_info(checkpoint_version_id) + civitai_info, _ = civitai_info_tuple if isinstance(civitai_info_tuple, tuple) else (civitai_info_tuple, None) + checkpoint = await self.populate_checkpoint_from_civitai(checkpoint, civitai_info) + except Exception as e: + logger.error(f"Error fetching Civitai info for checkpoint: {e}") recipe_base_model = checkpoint.get('baseModel') if checkpoint else None loras = [] diff --git a/py/services/batch_import_service.py b/py/services/batch_import_service.py index 7f558acc..7906ad42 100644 --- a/py/services/batch_import_service.py +++ b/py/services/batch_import_service.py @@ -118,6 +118,10 @@ class AdaptiveConcurrencyController: self._task_durations: List[float] = [] self._recent_errors = 0 self._recent_successes = 0 + # Batch-wide shared semaphore; created lazily on first use so the + # controller can also be constructed outside a running event loop. + self._semaphore: Optional[asyncio.Semaphore] = None + self._semaphore_capacity = initial_concurrency def record_result(self, duration: float, success: bool) -> None: self._task_durations.append(duration) @@ -146,7 +150,37 @@ class AdaptiveConcurrencyController: self._recent_successes = 0 def get_semaphore(self) -> asyncio.Semaphore: - return asyncio.Semaphore(self.current_concurrency) + """Return the batch-wide shared semaphore. + + The same semaphore instance is returned for every item of a batch so + the configured concurrency bounds are actually enforced. Previously a + fresh semaphore was created per call, letting every item run + concurrently and hammering remote metadata providers without any + limit. + """ + if self._semaphore is None: + self._semaphore = asyncio.Semaphore(self.current_concurrency) + self._semaphore_capacity = self.current_concurrency + return self._semaphore + + async def apply_concurrency(self) -> None: + """Synchronize the shared semaphore capacity with ``current_concurrency``. + + Call after ``record_result`` (once per completed item). Growing the + capacity is immediate (release). Shrinking requires acquiring a permit + and holding it, which is best-effort while other tasks are still + running — the capacity converges on subsequent calls. + """ + semaphore = self.get_semaphore() + while self._semaphore_capacity < self.current_concurrency: + semaphore.release() + self._semaphore_capacity += 1 + while self._semaphore_capacity > self.current_concurrency: + try: + await asyncio.wait_for(semaphore.acquire(), timeout=0.01) + except (asyncio.TimeoutError, asyncio.CancelledError): + break + self._semaphore_capacity -= 1 class BatchImportService: @@ -394,6 +428,9 @@ class BatchImportService: self._concurrency_controller.record_result( duration, result.get("success", False) ) + # Keep the shared batch semaphore in sync with the adaptively + # adjusted concurrency so the bounds actually take effect. + await self._concurrency_controller.apply_concurrency() if result.get("success"): item.status = ImportStatus.SUCCESS @@ -416,6 +453,7 @@ class BatchImportService: item.duration = time.time() - start_time progress.failed += 1 self._concurrency_controller.record_result(item.duration, False) + await self._concurrency_controller.apply_concurrency() progress.completed += 1 self._logger.info( diff --git a/py/services/civarchive_client.py b/py/services/civarchive_client.py index 0c300ea7..11de0d5f 100644 --- a/py/services/civarchive_client.py +++ b/py/services/civarchive_client.py @@ -7,6 +7,7 @@ import logging import asyncio from copy import deepcopy from typing import Any, Optional, Dict, Tuple, List, cast +from .connectivity_guard import is_expected_offline_error from .model_metadata_provider import CivArchiveModelMetadataProvider, ModelMetadataProviderManager from .downloader import get_downloader from .errors import RateLimitError @@ -46,7 +47,11 @@ class CivArchiveClient: """Call CivArchive API and return JSON payload""" success, payload = await self._make_request(path, params=params) if not success: - error = payload if isinstance(payload, str) else "Request failed" + # Normalize empty-string failure payloads (e.g. a throttled + # connection dropped without a message) so callers never see a + # falsy error alongside a None payload — that combination used to + # crash downstream None.get() calls. + error = payload if isinstance(payload, str) and payload else "Request failed" return None, error if not isinstance(payload, dict): return None, "Invalid response structure" @@ -298,6 +303,8 @@ class CivArchiveClient: async def _resolve_version_from_files(self, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]: """Fallback to fetch version data when only file metadata is available""" + if not isinstance(payload, dict): + return None data = self._normalize_payload(payload) files = data.get("files") or payload.get("files") or [] if not isinstance(files, list): @@ -332,10 +339,13 @@ class CivArchiveClient: """Find model by SHA256 hash value using CivArchive API""" try: payload, error = await self._request_json(f"/sha256/{model_hash.lower()}") - if error: - if "not found" in error.lower(): + # Treat a missing payload as an error even when the error string is + # falsy; passing None into the split/transform helpers below used to + # crash with "'NoneType' object has no attribute 'get'". + if error is not None or payload is None: + if error and "not found" in error.lower(): return None, "Model not found" - return None, error + return None, error or "Request failed" context, version_data, fallback_files = self._split_context(cast(Dict[str, Any], payload)) transformed = self._transform_version(context, version_data, fallback_files) @@ -352,7 +362,14 @@ class CivArchiveClient: except RateLimitError: raise except Exception as e: - logger.error(f"Error fetching CivArchive model by hash {model_hash[:10]}: {e}") + if is_expected_offline_error(str(e)): + logger.debug( + "Skipping CivArchive model by hash %s while offline: %s", + model_hash[:10], + e, + ) + else: + logger.error(f"Error fetching CivArchive model by hash {model_hash[:10]}: {e}") return None, str(e) async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]: @@ -362,7 +379,14 @@ class CivArchiveClient: if error or payload is None: if error and "not found" in error.lower(): return None - logger.error(f"Error fetching CivArchive model versions for {model_id}: {error}") + if is_expected_offline_error(error): + logger.debug( + "Skipping CivArchive model versions fetch for %s while offline: %s", + model_id, + error, + ) + else: + logger.error(f"Error fetching CivArchive model versions for {model_id}: {error}") return None data = self._normalize_payload(payload) @@ -426,7 +450,19 @@ class CivArchiveClient: if error or payload is None: if error and "not found" in error.lower(): return None - logger.error(f"Error fetching CivArchive model version via API {model_id}/{version_id}: {error}") + # The connectivity guard short-circuits requests during its + # offline cooldown; that is an expected, transient state, so + # log it as DEBUG instead of spamming one ERROR per request + # (batch imports can hit this thousands of times). + if is_expected_offline_error(error): + logger.debug( + "Skipping CivArchive model version fetch %s/%s while offline: %s", + model_id, + version_id, + error, + ) + else: + logger.error(f"Error fetching CivArchive model version via API {model_id}/{version_id}: {error}") return None context, version_data, fallback_files = self._split_context(payload) diff --git a/tests/services/test_batch_import_service.py b/tests/services/test_batch_import_service.py index 89e41f8e..3cd8e4f7 100644 --- a/tests/services/test_batch_import_service.py +++ b/tests/services/test_batch_import_service.py @@ -154,6 +154,65 @@ class TestAdaptiveConcurrencyController: controller.record_result(duration=5.0, success=True) assert controller.current_concurrency == 3 + @pytest.mark.asyncio + async def test_get_semaphore_returns_shared_instance(self): + controller = AdaptiveConcurrencyController(initial_concurrency=3) + # Every item of a batch must receive the same semaphore so the + # concurrency bound is actually enforced batch-wide. + first = controller.get_semaphore() + second = controller.get_semaphore() + assert first is second + + @pytest.mark.asyncio + async def test_shared_semaphore_limits_concurrent_tasks(self): + controller = AdaptiveConcurrencyController(initial_concurrency=3) + semaphore = controller.get_semaphore() + active = 0 + peak = 0 + + async def worker(): + nonlocal active, peak + async with semaphore: + active += 1 + peak = max(peak, active) + await asyncio.sleep(0.05) + active -= 1 + + await asyncio.gather(*[worker() for _ in range(10)]) + assert peak == 3 + + @pytest.mark.asyncio + async def test_apply_concurrency_increases_capacity(self): + controller = AdaptiveConcurrencyController(initial_concurrency=3) + semaphore = controller.get_semaphore() + controller.record_result(duration=0.5, success=True) # 3 -> 4 + await controller.apply_concurrency() + + acquired = await asyncio.gather( + *[asyncio.wait_for(semaphore.acquire(), timeout=0.2) for _ in range(4)] + ) + assert all(acquired) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(semaphore.acquire(), timeout=0.05) + for _ in range(4): + semaphore.release() + + @pytest.mark.asyncio + async def test_apply_concurrency_decreases_capacity(self): + controller = AdaptiveConcurrencyController(initial_concurrency=3) + semaphore = controller.get_semaphore() + controller.record_result(duration=1.0, success=False) # 3 -> 2 + await controller.apply_concurrency() + + acquired = await asyncio.gather( + *[asyncio.wait_for(semaphore.acquire(), timeout=0.2) for _ in range(2)] + ) + assert all(acquired) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(semaphore.acquire(), timeout=0.05) + for _ in range(2): + semaphore.release() + class TestBatchImportProgress: def test_to_dict(self): diff --git a/tests/services/test_civarchive_client.py b/tests/services/test_civarchive_client.py index 507e1094..5eaca7bf 100644 --- a/tests/services/test_civarchive_client.py +++ b/tests/services/test_civarchive_client.py @@ -1,4 +1,5 @@ import copy +import logging from typing import Any, Dict from unittest.mock import AsyncMock @@ -257,3 +258,49 @@ async def test_get_model_by_hash_propagates_rate_limit(downloader): assert exc_info.value.retry_after == 5 assert exc_info.value.provider == "civarchive_api" + + +async def test_get_model_by_hash_empty_error_payload(downloader): + """An empty-string failure payload must surface as a proper error. + + Regression test: (None, "") used to fall through the falsy-error check and + crash in _resolve_version_from_files with "'NoneType' object has no + attribute 'get'". + """ + + async def fake_make_request(method, url, use_auth=False, **kwargs): + return False, "" + + downloader.make_request = fake_make_request + + client = await CivArchiveClient.get_instance() + + result, error = await client.get_model_by_hash("empty-error") + + assert result is None + assert error == "Request failed" + + +async def test_get_model_version_offline_cooldown_logged_as_debug(downloader, caplog): + """Cooldown short-circuits must not spam an ERROR per request.""" + + async def fake_make_request(method, url, use_auth=False, **kwargs): + return False, "offline_cooldown" + + downloader.make_request = fake_make_request + + client = await CivArchiveClient.get_instance() + + with caplog.at_level(logging.DEBUG, logger="py.services.civarchive_client"): + result = await client.get_model_version(model_id=1, version_id=123) + + assert result is None + module_records = [r for r in caplog.records if r.name == "py.services.civarchive_client"] + assert not any( + r.levelno >= logging.ERROR and "Error fetching CivArchive model version" in r.getMessage() + for r in module_records + ) + assert any( + r.levelno == logging.DEBUG and "while offline" in r.getMessage() + for r in module_records + ) diff --git a/tests/services/test_comfy_metadata_parser.py b/tests/services/test_comfy_metadata_parser.py index 25136f56..75266403 100644 --- a/tests/services/test_comfy_metadata_parser.py +++ b/tests/services/test_comfy_metadata_parser.py @@ -277,3 +277,72 @@ async def test_parse_metadata_without_extra_metadata(monkeypatch): assert "error" not in result assert result["loras"] == [] assert result["checkpoint"]["id"] == "456" + + +@pytest.mark.asyncio +async def test_parse_metadata_with_list_ckpt_name(monkeypatch): + """ckpt_name serialized as a single-element list must not crash re.search.""" + checkpoint_info = { + "id": 456, + "modelId": 123, + "model": {"name": "Checkpoint", "type": "checkpoint"}, + "name": "v1", + "baseModel": "SDXL 1.0", + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + assert version_id == "456" + return checkpoint_info, None + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.comfy.get_default_metadata_provider", + fake_metadata_provider, + ) + + metadata_json = { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": ["urn:air:sdxl:checkpoint:civitai:123@456"]}, + } + } + + result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json)) + + assert "error" not in result + assert result["checkpoint"] is not None + assert int(result["checkpoint"]["id"]) == 456 + assert int(result["checkpoint"]["modelId"]) == 123 + + +@pytest.mark.asyncio +async def test_parse_metadata_with_none_ckpt_name(monkeypatch): + """Missing (None) ckpt_name must not crash re.search with a TypeError.""" + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + raise AssertionError("Checkpoint lookup must be skipped") + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.comfy.get_default_metadata_provider", + fake_metadata_provider, + ) + + metadata_json = { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": None}, + } + } + + result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json)) + + assert "error" not in result + assert result["checkpoint"] is None + assert result["loras"] == []