mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
Compare commits
3 Commits
v1.2.2
...
a7995db009
| Author | SHA1 | Date | |
|---|---|---|---|
| a7995db009 | |||
| 5ae4aef30e | |||
| 08023f0cd9 |
+91
-48
@@ -11,6 +11,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
@@ -32,8 +33,26 @@ _catalog_cache: Optional[Dict[str, List[str]]] = None
|
||||
# ``{provider_id: {model_id: max_output_tokens}}``.
|
||||
_model_output_limits: Dict[str, Dict[str, int]] = {}
|
||||
|
||||
# Monotonic timestamp of the last failed catalog fetch (None = no failure
|
||||
# yet). Failed fetches are negatively cached: further calls return the
|
||||
# empty fallback without hitting the network until the cooldown elapses,
|
||||
# so users on broken networks don't stall on every settings-modal open.
|
||||
_catalog_last_failure: Optional[float] = None
|
||||
_CATALOG_FAILURE_COOLDOWN = 600.0 # seconds
|
||||
|
||||
# Serializes catalog fetches so concurrent callers don't duplicate requests.
|
||||
_catalog_lock = asyncio.Lock()
|
||||
|
||||
_CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30)
|
||||
|
||||
# Cloudflare serves brotli when the client advertises it, and brotli is a
|
||||
# required dependency here — a corrupted br stream can crash the native
|
||||
# decoder with a Windows access violation (issue #1099). Request gzip
|
||||
# instead; zlib decompression is not affected and corrupt gzip data only
|
||||
# raises ContentEncodingError (an aiohttp.ClientError subclass), which the
|
||||
# exception handlers below already catch.
|
||||
_NO_BROTLI_HEADERS = {"Accept-Encoding": "gzip, deflate"}
|
||||
|
||||
|
||||
async def _load_model_catalog() -> Dict[str, List[str]]:
|
||||
"""Fetch and parse the model catalog.
|
||||
@@ -46,61 +65,85 @@ async def _load_model_catalog() -> Dict[str, List[str]]:
|
||||
value has a ``models`` sub-dict keyed by model ID. The result is cached
|
||||
in memory after the first successful fetch.
|
||||
Subsequent calls return the cached data immediately.
|
||||
|
||||
Failed fetches are negatively cached: further calls return an empty
|
||||
dict without hitting the network until ``_CATALOG_FAILURE_COOLDOWN``
|
||||
has elapsed, so a broken network does not stall every settings-modal
|
||||
open. Concurrent callers are serialized behind :data:`_catalog_lock`
|
||||
so only one request is ever in flight.
|
||||
"""
|
||||
global _catalog_cache, _model_output_limits
|
||||
global _catalog_cache, _model_output_limits, _catalog_last_failure
|
||||
if _catalog_cache is not None:
|
||||
return _catalog_cache
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session:
|
||||
async with session.get(_MODEL_CATALOG_URL) as resp:
|
||||
if resp.status != 200:
|
||||
logger.warning("Model catalog returned HTTP %s", resp.status)
|
||||
return _catalog_cache or {}
|
||||
data = await resp.json()
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc:
|
||||
logger.warning("Failed to fetch model catalog: %s", exc)
|
||||
return _catalog_cache or {}
|
||||
async with _catalog_lock:
|
||||
# Re-check under the lock: another caller may have fetched (or
|
||||
# failed) while we were waiting.
|
||||
if _catalog_cache is not None:
|
||||
return _catalog_cache
|
||||
if (
|
||||
_catalog_last_failure is not None
|
||||
and time.monotonic() - _catalog_last_failure < _CATALOG_FAILURE_COOLDOWN
|
||||
):
|
||||
logger.debug(
|
||||
"Skipping model catalog fetch: last attempt failed %.0fs ago",
|
||||
time.monotonic() - _catalog_last_failure,
|
||||
)
|
||||
return {}
|
||||
|
||||
if not isinstance(data, dict):
|
||||
logger.warning("Model catalog is not a dict, got %s", type(data).__name__)
|
||||
return _catalog_cache or {}
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session:
|
||||
async with session.get(_MODEL_CATALOG_URL, headers=_NO_BROTLI_HEADERS) as resp:
|
||||
if resp.status != 200:
|
||||
logger.warning("Model catalog returned HTTP %s", resp.status)
|
||||
_catalog_last_failure = time.monotonic()
|
||||
return {}
|
||||
data = await resp.json()
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc:
|
||||
logger.warning("Failed to fetch model catalog: %s", exc)
|
||||
_catalog_last_failure = time.monotonic()
|
||||
return {}
|
||||
|
||||
result: Dict[str, List[str]] = {}
|
||||
output_limits: Dict[str, Dict[str, int]] = {}
|
||||
for provider_id, provider_info in data.items():
|
||||
if not isinstance(provider_info, dict):
|
||||
continue
|
||||
models_dict = provider_info.get("models")
|
||||
if not isinstance(models_dict, dict):
|
||||
continue
|
||||
model_ids: List[str] = []
|
||||
provider_limits: Dict[str, int] = {}
|
||||
for mid, model_info in models_dict.items():
|
||||
if not isinstance(mid, str):
|
||||
if not isinstance(data, dict):
|
||||
logger.warning("Model catalog is not a dict, got %s", type(data).__name__)
|
||||
_catalog_last_failure = time.monotonic()
|
||||
return {}
|
||||
|
||||
result: Dict[str, List[str]] = {}
|
||||
output_limits: Dict[str, Dict[str, int]] = {}
|
||||
for provider_id, provider_info in data.items():
|
||||
if not isinstance(provider_info, dict):
|
||||
continue
|
||||
model_ids.append(mid)
|
||||
if isinstance(model_info, dict):
|
||||
limit = model_info.get("limit")
|
||||
if isinstance(limit, dict):
|
||||
output = limit.get("output")
|
||||
if isinstance(output, (int, float)) and output > 0:
|
||||
provider_limits[mid] = int(output)
|
||||
if model_ids:
|
||||
result[provider_id] = model_ids
|
||||
if provider_limits:
|
||||
output_limits[provider_id] = provider_limits
|
||||
models_dict = provider_info.get("models")
|
||||
if not isinstance(models_dict, dict):
|
||||
continue
|
||||
model_ids: List[str] = []
|
||||
provider_limits: Dict[str, int] = {}
|
||||
for mid, model_info in models_dict.items():
|
||||
if not isinstance(mid, str):
|
||||
continue
|
||||
model_ids.append(mid)
|
||||
if isinstance(model_info, dict):
|
||||
limit = model_info.get("limit")
|
||||
if isinstance(limit, dict):
|
||||
output = limit.get("output")
|
||||
if isinstance(output, (int, float)) and output > 0:
|
||||
provider_limits[mid] = int(output)
|
||||
if model_ids:
|
||||
result[provider_id] = model_ids
|
||||
if provider_limits:
|
||||
output_limits[provider_id] = provider_limits
|
||||
|
||||
_catalog_cache = result
|
||||
_model_output_limits = output_limits
|
||||
logger.debug(
|
||||
"Loaded model catalog: %d providers, %d total models "
|
||||
"(%d providers have output limits)",
|
||||
len(result),
|
||||
sum(len(m) for m in result.values()),
|
||||
len(output_limits),
|
||||
)
|
||||
return result
|
||||
_catalog_cache = result
|
||||
_model_output_limits = output_limits
|
||||
logger.debug(
|
||||
"Loaded model catalog: %d providers, %d total models "
|
||||
"(%d providers have output limits)",
|
||||
len(result),
|
||||
sum(len(m) for m in result.values()),
|
||||
len(output_limits),
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _get_model_max_output(provider: str, model: str) -> Optional[int]:
|
||||
@@ -126,7 +169,7 @@ async def fetch_ollama_models(api_base: str) -> List[str]:
|
||||
url = f"{api_base.rstrip('/')}/models"
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=_OLLAMA_API_TIMEOUT) as session:
|
||||
async with session.get(url) as resp:
|
||||
async with session.get(url, headers=_NO_BROTLI_HEADERS) as resp:
|
||||
if resp.status != 200:
|
||||
logger.debug("Ollama API returned HTTP %s from %s", resp.status, api_base)
|
||||
return []
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
@@ -322,8 +323,12 @@ class MockGetSession:
|
||||
|
||||
def __init__(self, response):
|
||||
self._response = response
|
||||
self.last_url = None
|
||||
self.last_headers = None
|
||||
|
||||
def get(self, url):
|
||||
def get(self, url, headers=None):
|
||||
self.last_url = url
|
||||
self.last_headers = headers
|
||||
return self._response
|
||||
|
||||
async def __aenter__(self):
|
||||
@@ -340,6 +345,14 @@ class CorruptJsonResponse(MockResponse):
|
||||
raise UnicodeDecodeError("utf-8", b"\x9a", 0, 1, "invalid start byte")
|
||||
|
||||
|
||||
class SlowResponse(MockResponse):
|
||||
"""Response whose body takes a moment to read, to force contention."""
|
||||
|
||||
async def json(self):
|
||||
await asyncio.sleep(0.05)
|
||||
return self._json_data
|
||||
|
||||
|
||||
class TestModelCatalog:
|
||||
"""Tests for _load_model_catalog / fetch_ollama_models error handling."""
|
||||
|
||||
@@ -348,9 +361,11 @@ class TestModelCatalog:
|
||||
"""Reset the module-level catalog cache around each test."""
|
||||
llm_module._catalog_cache = None
|
||||
llm_module._model_output_limits = {}
|
||||
llm_module._catalog_last_failure = None
|
||||
yield
|
||||
llm_module._catalog_cache = None
|
||||
llm_module._model_output_limits = {}
|
||||
llm_module._catalog_last_failure = None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_model_catalog_falls_back_on_unicode_decode_error(self):
|
||||
@@ -373,3 +388,86 @@ class TestModelCatalog:
|
||||
models = await fetch_ollama_models("http://localhost:11434/v1")
|
||||
|
||||
assert models == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_catalog_request_disables_brotli_encoding(self):
|
||||
"""The catalog request must not advertise br — a corrupt brotli stream
|
||||
can crash the native decoder (Windows access violation, issue #1099)."""
|
||||
response = MockResponse(200, json_data={})
|
||||
session = MockGetSession(response)
|
||||
|
||||
with mock.patch("aiohttp.ClientSession", return_value=session):
|
||||
await llm_module._load_model_catalog()
|
||||
|
||||
assert session.last_headers == {"Accept-Encoding": "gzip, deflate"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_request_disables_brotli_encoding(self):
|
||||
"""The Ollama models request must not advertise br either."""
|
||||
response = MockResponse(200, json_data={"data": [{"id": "llama3"}]})
|
||||
session = MockGetSession(response)
|
||||
|
||||
with mock.patch("aiohttp.ClientSession", return_value=session):
|
||||
models = await fetch_ollama_models("http://localhost:11434/v1")
|
||||
|
||||
assert models == ["llama3"]
|
||||
assert session.last_headers == {"Accept-Encoding": "gzip, deflate"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_fetch_is_negatively_cached(self):
|
||||
"""A failed fetch is not retried until the cooldown elapses."""
|
||||
created = []
|
||||
|
||||
def factory(*args, **kwargs):
|
||||
session = MockGetSession(MockResponse(500, text_data="error"))
|
||||
created.append(session)
|
||||
return session
|
||||
|
||||
with mock.patch("aiohttp.ClientSession", side_effect=factory):
|
||||
first = await llm_module._load_model_catalog()
|
||||
second = await llm_module._load_model_catalog()
|
||||
|
||||
assert first == {}
|
||||
assert second == {}
|
||||
assert len(created) == 1
|
||||
assert llm_module._catalog_last_failure is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_retries_after_cooldown(self):
|
||||
"""Once the cooldown elapses, the next call fetches again."""
|
||||
bad = MockGetSession(MockResponse(500, text_data="error"))
|
||||
with mock.patch("aiohttp.ClientSession", return_value=bad):
|
||||
assert await llm_module._load_model_catalog() == {}
|
||||
|
||||
# Simulate the cooldown having elapsed.
|
||||
llm_module._catalog_last_failure = (
|
||||
time.monotonic() - llm_module._CATALOG_FAILURE_COOLDOWN - 1
|
||||
)
|
||||
|
||||
good = MockGetSession(
|
||||
MockResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}})
|
||||
)
|
||||
with mock.patch("aiohttp.ClientSession", return_value=good):
|
||||
catalog = await llm_module._load_model_catalog()
|
||||
|
||||
assert catalog == {"openai": ["gpt-4o"]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_fetches_are_deduplicated(self):
|
||||
"""Concurrent callers share a single in-flight fetch."""
|
||||
created = []
|
||||
|
||||
def factory(*args, **kwargs):
|
||||
session = MockGetSession(
|
||||
SlowResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}})
|
||||
)
|
||||
created.append(session)
|
||||
return session
|
||||
|
||||
with mock.patch("aiohttp.ClientSession", side_effect=factory):
|
||||
results = await asyncio.gather(
|
||||
*(llm_module._load_model_catalog() for _ in range(3))
|
||||
)
|
||||
|
||||
assert len(created) == 1
|
||||
assert all(r == {"openai": ["gpt-4o"]} for r in results)
|
||||
|
||||
Reference in New Issue
Block a user