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
@@ -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):
+47
View File
@@ -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"] == []