mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-28 00:11:27 -03:00
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:
@@ -40,6 +40,16 @@ 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']
|
||||
# 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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,6 +362,13 @@ class CivArchiveClient:
|
||||
except RateLimitError:
|
||||
raise
|
||||
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}")
|
||||
return None, str(e)
|
||||
|
||||
@@ -362,6 +379,13 @@ class CivArchiveClient:
|
||||
if error or payload is None:
|
||||
if error and "not found" in error.lower():
|
||||
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}")
|
||||
return None
|
||||
|
||||
@@ -426,6 +450,18 @@ class CivArchiveClient:
|
||||
if error or payload is None:
|
||||
if error and "not found" in error.lower():
|
||||
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}")
|
||||
return None
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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"] == []
|
||||
|
||||
Reference in New Issue
Block a user