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.
This commit is contained in:
Will Miao
2026-08-27 07:58:48 +08:00
parent 574dfbbe55
commit ee233548e5
6 changed files with 285 additions and 26 deletions
+10
View File
@@ -40,6 +40,16 @@ class ComfyMetadataParser(RecipeMetadataParser):
checkpoint_node = next(iter(checkpoint_nodes.values())) checkpoint_node = next(iter(checkpoint_nodes.values()))
if 'inputs' in checkpoint_node and 'ckpt_name' in checkpoint_node['inputs']: if 'inputs' in checkpoint_node and 'ckpt_name' in checkpoint_node['inputs']:
checkpoint_name = checkpoint_node['inputs']['ckpt_name'] checkpoint_name = checkpoint_node['inputs']['ckpt_name']
# 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) checkpoint_match = re.search(r'civitai:(\d+)@(\d+)', checkpoint_name)
if checkpoint_match: if checkpoint_match:
checkpoint_id = checkpoint_match.group(1) checkpoint_id = checkpoint_match.group(1)
+39 -1
View File
@@ -118,6 +118,10 @@ class AdaptiveConcurrencyController:
self._task_durations: List[float] = [] self._task_durations: List[float] = []
self._recent_errors = 0 self._recent_errors = 0
self._recent_successes = 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: def record_result(self, duration: float, success: bool) -> None:
self._task_durations.append(duration) self._task_durations.append(duration)
@@ -146,7 +150,37 @@ class AdaptiveConcurrencyController:
self._recent_successes = 0 self._recent_successes = 0
def get_semaphore(self) -> asyncio.Semaphore: 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: class BatchImportService:
@@ -394,6 +428,9 @@ class BatchImportService:
self._concurrency_controller.record_result( self._concurrency_controller.record_result(
duration, result.get("success", False) 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"): if result.get("success"):
item.status = ImportStatus.SUCCESS item.status = ImportStatus.SUCCESS
@@ -416,6 +453,7 @@ class BatchImportService:
item.duration = time.time() - start_time item.duration = time.time() - start_time
progress.failed += 1 progress.failed += 1
self._concurrency_controller.record_result(item.duration, False) self._concurrency_controller.record_result(item.duration, False)
await self._concurrency_controller.apply_concurrency()
progress.completed += 1 progress.completed += 1
self._logger.info( self._logger.info(
+40 -4
View File
@@ -7,6 +7,7 @@ import logging
import asyncio import asyncio
from copy import deepcopy from copy import deepcopy
from typing import Any, Optional, Dict, Tuple, List, cast 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 .model_metadata_provider import CivArchiveModelMetadataProvider, ModelMetadataProviderManager
from .downloader import get_downloader from .downloader import get_downloader
from .errors import RateLimitError from .errors import RateLimitError
@@ -46,7 +47,11 @@ class CivArchiveClient:
"""Call CivArchive API and return JSON payload""" """Call CivArchive API and return JSON payload"""
success, payload = await self._make_request(path, params=params) success, payload = await self._make_request(path, params=params)
if not success: 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 return None, error
if not isinstance(payload, dict): if not isinstance(payload, dict):
return None, "Invalid response structure" 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]]: 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""" """Fallback to fetch version data when only file metadata is available"""
if not isinstance(payload, dict):
return None
data = self._normalize_payload(payload) data = self._normalize_payload(payload)
files = data.get("files") or payload.get("files") or [] files = data.get("files") or payload.get("files") or []
if not isinstance(files, list): if not isinstance(files, list):
@@ -332,10 +339,13 @@ class CivArchiveClient:
"""Find model by SHA256 hash value using CivArchive API""" """Find model by SHA256 hash value using CivArchive API"""
try: try:
payload, error = await self._request_json(f"/sha256/{model_hash.lower()}") payload, error = await self._request_json(f"/sha256/{model_hash.lower()}")
if error: # Treat a missing payload as an error even when the error string is
if "not found" in error.lower(): # 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, "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)) context, version_data, fallback_files = self._split_context(cast(Dict[str, Any], payload))
transformed = self._transform_version(context, version_data, fallback_files) transformed = self._transform_version(context, version_data, fallback_files)
@@ -352,6 +362,13 @@ class CivArchiveClient:
except RateLimitError: except RateLimitError:
raise raise
except Exception as e: except Exception as 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}") logger.error(f"Error fetching CivArchive model by hash {model_hash[:10]}: {e}")
return None, str(e) return None, str(e)
@@ -362,6 +379,13 @@ class CivArchiveClient:
if error or payload is None: if error or payload is None:
if error and "not found" in error.lower(): if error and "not found" in error.lower():
return None return None
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}") logger.error(f"Error fetching CivArchive model versions for {model_id}: {error}")
return None return None
@@ -426,6 +450,18 @@ class CivArchiveClient:
if error or payload is None: if error or payload is None:
if error and "not found" in error.lower(): if error and "not found" in error.lower():
return None return None
# 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}") logger.error(f"Error fetching CivArchive model version via API {model_id}/{version_id}: {error}")
return None return None
@@ -154,6 +154,65 @@ class TestAdaptiveConcurrencyController:
controller.record_result(duration=5.0, success=True) controller.record_result(duration=5.0, success=True)
assert controller.current_concurrency == 3 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: class TestBatchImportProgress:
def test_to_dict(self): def test_to_dict(self):
+47
View File
@@ -1,4 +1,5 @@
import copy import copy
import logging
from typing import Any, Dict from typing import Any, Dict
from unittest.mock import AsyncMock 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.retry_after == 5
assert exc_info.value.provider == "civarchive_api" 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
)
@@ -277,3 +277,72 @@ async def test_parse_metadata_without_extra_metadata(monkeypatch):
assert "error" not in result assert "error" not in result
assert result["loras"] == [] assert result["loras"] == []
assert result["checkpoint"]["id"] == "456" 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"] == []