From e9e9ee20c6a2ba58c09da749e84182f1224e8073 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 14 Sep 2026 21:24:57 +0800 Subject: [PATCH] perf(modelscope): fetch a repository's model card once per run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A collection repository publishes many model files under a single source id, but enrichment re-read the README and the model-detail payload for every one of them: eight checkpoints meant sixteen HTTP requests, each detail payload being 10-22 KB of JSON. Add `ModelSourceCache`, created by `execute_skill()` for the duration of a run and passed to the provider through a new optional `cache` argument on `fetch_model_card_context()`. The agent caches the README (repository-wide and provider-agnostic), and ModelScope caches its detail payload under a provider-namespaced key. Only successful reads are memoised, so a transient failure is still retried for the next file, and the per-file selection is redone from the cached payload so a checkpoint never inherits a sibling's example images. Nothing is retained across runs — a model card can change at any time — and download URLs are not routed through the cache. Measured over the eight checkpoints of one ModelScope repository: 16 requests before, 2 after. To keep the two concerns separable, `_build_card_context()` now turns a detail payload into a `ModelCardContext` as a pure function. --- docs/agent_skills.md | 7 ++ py/services/agent/agent_service.py | 30 ++++- py/services/model_sources/__init__.py | 2 + py/services/model_sources/base.py | 33 +++++- py/services/model_sources/modelscope.py | 81 ++++++++++--- .../services/test_agent_enrichment_source.py | 9 +- tests/services/test_model_sources.py | 111 ++++++++++++++++++ 7 files changed, 247 insertions(+), 26 deletions(-) diff --git a/docs/agent_skills.md b/docs/agent_skills.md index bcc264fd..c45152d2 100644 --- a/docs/agent_skills.md +++ b/docs/agent_skills.md @@ -109,6 +109,13 @@ returns a `ModelCardContext`. Example images are matched to the model's Sites with no such extras inherit an empty context, and the pipeline behaves exactly as before. +The README and the repository metadata describe the whole repository, not one +file, so `execute_skill()` creates a `ModelSourceCache` for the duration of a +run and passes it down. Enriching the eight checkpoints of one ModelScope +repository costs two HTTP requests instead of sixteen; only the per-file +selection is redone for each file. Nothing is cached across runs, and download +URLs never go through it. + #### Deterministic data is applied whether or not an LLM is configured `AgentService._load_source_card()` runs for every source-backed enrichment, and diff --git a/py/services/agent/agent_service.py b/py/services/agent/agent_service.py index 1440de0c..10cf3461 100644 --- a/py/services/agent/agent_service.py +++ b/py/services/agent/agent_service.py @@ -28,6 +28,7 @@ from ...config import config from ..llm_service import LLMService from ..model_sources import ( ModelCardContext, + ModelSourceCache, get_source, resolve_source_ref, source_label, @@ -259,6 +260,11 @@ class AgentService: llm = await self._ensure_llm() llm_configured = llm.is_configured() if skill.llm_required else True + # A collection repository holds many model files under one source id; + # this memo keeps the README and the repository metadata from being + # re-fetched once per file. It lives for this run only. + source_cache = ModelSourceCache() + for model_path in model_paths: model_filename = os.path.basename(model_path) logger.info( @@ -288,7 +294,7 @@ class AgentService: # or not an LLM is available: a user without a key still gets # the author summary, the example images and the tags. source_vars, source_context = await self._load_source_card( - model_path, metadata, + model_path, metadata, cache=source_cache, ) resolved_base_model = "" if skill_name == "enrich_hf_metadata" and not ( @@ -423,13 +429,22 @@ class AgentService: return "\n".join(f"- {m}" for m in models) async def _load_source_card( - self, model_path: str, metadata: Dict[str, Any] + self, + model_path: str, + metadata: Dict[str, Any], + *, + cache: Optional[ModelSourceCache] = None, ) -> tuple[Dict[str, Any], ModelCardContext]: """Fetch the model card and site-published extras for one model. Runs for every source-backed enrichment regardless of LLM availability, because everything it returns is deterministic data that should be applied even without a configured provider. + + *cache* is the per-run memo created by :meth:`execute_skill`. The + README is repository-wide, so it is fetched once per source id; only + successful reads are memoised, leaving a transient failure to be + retried for the next file. """ variables: Dict[str, Any] = { @@ -450,11 +465,18 @@ class AgentService: raw_basename = os.path.splitext(os.path.basename(model_path))[0] variables["asset_base_url"] = source.asset_base_url(ref.source_id) - readme = await source.fetch_model_card(ref.source_id) + + cache_key = f"{ref.platform}:{ref.source_id}" + readme = cache.readmes.get(cache_key) if cache is not None else None + if readme is None: + readme = await source.fetch_model_card(ref.source_id) + if cache is not None and readme: + cache.readmes[cache_key] = readme + # Sites such as ModelScope keep part of the model card outside the # README (author summary, curated tags, per-file example images). card_context = await source.fetch_model_card_context( - ref.source_id, os.path.basename(model_path) + ref.source_id, os.path.basename(model_path), cache=cache, ) variables["source_description"] = card_context.description variables["source_base_model"] = card_context.base_model diff --git a/py/services/model_sources/__init__.py b/py/services/model_sources/__init__.py index 290b7ecc..555795fb 100644 --- a/py/services/model_sources/__init__.py +++ b/py/services/model_sources/__init__.py @@ -13,6 +13,7 @@ from .base import ( HTTP_TIMEOUT, ModelCardContext, ModelSource, + ModelSourceCache, ModelSourceError, SourceRef, USER_AGENT, @@ -48,6 +49,7 @@ __all__ = [ "LEGACY_HF_URL_FIELD", "ModelCardContext", "ModelSource", + "ModelSourceCache", "ModelSourceError", "HuggingFaceSource", "ModelScopeSource", diff --git a/py/services/model_sources/base.py b/py/services/model_sources/base.py index b0b13a45..49f63bd9 100644 --- a/py/services/model_sources/base.py +++ b/py/services/model_sources/base.py @@ -25,7 +25,7 @@ import logging import os import re from dataclasses import dataclass, field -from typing import Any, Iterable, Optional +from typing import Any, Dict, Iterable, Optional import aiohttp @@ -125,6 +125,26 @@ class ModelSourceError(Exception): self.status = status +class ModelSourceCache: + """Per-run memo shared between the agent pipeline and a model source. + + A collection repository publishes many model files under a single source + id, so enriching each file re-fetches the same README and the same + repository metadata. One cache is created per enrichment run and thrown + away afterwards: nothing is retained across runs (a model card can change + at any time), and download URLs are never routed through it. + """ + + def __init__(self) -> None: + #: Provider-agnostic: ``":"`` → raw README text. + self.readmes: Dict[str, str] = {} + #: Provider-owned scratch space. Keys must be namespaced by the + #: provider (``(platform, kind, source_id)``) so two providers can + #: never collide. Only successful results should be stored, so a + #: transient failure is still retried for the next file. + self.provider: Dict[Any, Any] = {} + + #: Repository ids are always exactly ``owner/name``. Components may contain #: dots (``black-forest-labs/FLUX.1-dev``) but must not be empty, ``.`` / ``..``, #: or start with a dot - the id is used as a path segment on disk. @@ -283,7 +303,11 @@ class ModelSource: return "" async def fetch_model_card_context( - self, source_id: str, filename: str = "" + self, + source_id: str, + filename: str = "", + *, + cache: Optional["ModelSourceCache"] = None, ) -> ModelCardContext: """Return the card extras the site keeps outside the README. @@ -292,6 +316,10 @@ class ModelSource: model card is fully described by :meth:`fetch_model_card` need no override and inherit this empty context. + *cache* is an optional per-run memo (see :class:`ModelSourceCache`) + that lets a provider avoid re-fetching repository-wide data for every + file in a collection repository. + Implementations must never raise: enrichment treats a missing context as "the site had nothing extra to say". """ @@ -371,6 +399,7 @@ __all__ = [ "HTTP_TIMEOUT", "ModelCardContext", "ModelSource", + "ModelSourceCache", "ModelSourceError", "SourceRef", "USER_AGENT", diff --git a/py/services/model_sources/modelscope.py b/py/services/model_sources/modelscope.py index d0324d6c..f8e53008 100644 --- a/py/services/model_sources/modelscope.py +++ b/py/services/model_sources/modelscope.py @@ -24,6 +24,10 @@ which redirects to a CDN URL carrying a time-limited ``auth_key``. Requesting the resolve URL fresh on every attempt (which the shared downloader does, including for resumable Range requests) keeps that key valid; the CDN URL must never be cached. + +The README and the detail payload both describe the whole repository rather +than one file, so a per-run ``ModelSourceCache`` keeps them from being read +again for every checkpoint of a collection repository. """ from __future__ import annotations @@ -32,7 +36,7 @@ import json import logging import os import re -from typing import Any +from typing import TYPE_CHECKING, Any, Optional from .base import ( ModelCardContext, @@ -43,6 +47,9 @@ from .base import ( filter_weight_files, ) +if TYPE_CHECKING: # pragma: no cover - typing only + from .base import ModelSourceCache + logger = logging.getLogger(__name__) _URL_PATTERN = re.compile( @@ -106,7 +113,11 @@ class ModelScopeSource(ModelSource): return "" async def fetch_model_card_context( - self, source_id: str, filename: str = "" + self, + source_id: str, + filename: str = "", + *, + cache: Optional["ModelSourceCache"] = None, ) -> ModelCardContext: """Read the model-detail API that backs the ModelScope model page. @@ -121,30 +132,44 @@ class ModelScopeSource(ModelSource): ``stats.fileList``, which means the images returned belong to the exact ``.safetensors`` being enriched — essential for collection repositories, where every checkpoint has its own sample image. + + The detail payload describes the whole repository and is therefore + shared across every file in it, so it is read through *cache* when the + caller supplies one; only the per-file selection is redone. """ + data = await self._fetch_detail(source_id, cache=cache) + if data is None: + return ModelCardContext() + return _build_card_context(data, filename) + + async def _fetch_detail( + self, + source_id: str, + *, + cache: Optional["ModelSourceCache"] = None, + ) -> Optional[dict[str, Any]]: + """Fetch (or reuse) the model-detail payload for *source_id*.""" + + cache_key = (self.platform, "detail", source_id) + if cache is not None and cache_key in cache.provider: + return cache.provider[cache_key] + status, payload = await fetch_json( f"https://modelscope.cn/api/v1/models/{source_id}" ) if status != 200 or not isinstance(payload, dict): - logger.debug("ModelScope detail API returned HTTP %s for %s", status, source_id) - return ModelCardContext() + logger.debug( + "ModelScope detail API returned HTTP %s for %s", status, source_id + ) + return None data = payload.get("Data") if not isinstance(data, dict): - return ModelCardContext() + return None - context = ModelCardContext( - description=_clean_text(data.get("Description")), - base_model=_first_string(data.get("BaseModel")), - base_model_aliases=_base_model_aliases(data), - official_tags=_official_tags(data.get("OfficialTags")), - ) - - versions = _matching_versions(data.get("MuseInfo"), filename) - if versions: - context.example_images = _cover_image_urls(versions) - context.trigger_words = _version_trigger_words(versions) - return context + if cache is not None: + cache.provider[cache_key] = data + return data async def list_files( self, source_id: str, revision: str = "" @@ -220,6 +245,28 @@ def _first_string(value: Any) -> str: return "" +def _build_card_context(data: dict[str, Any], filename: str) -> ModelCardContext: + """Turn a model-detail payload into a :class:`ModelCardContext`. + + Separated from the HTTP fetch so the repository-wide payload can be cached + across the files of a collection repository while the per-file selection + is still redone for each one. + """ + + context = ModelCardContext( + description=_clean_text(data.get("Description")), + base_model=_first_string(data.get("BaseModel")), + base_model_aliases=_base_model_aliases(data), + official_tags=_official_tags(data.get("OfficialTags")), + ) + + versions = _matching_versions(data.get("MuseInfo"), filename) + if versions: + context.example_images = _cover_image_urls(versions) + context.trigger_words = _version_trigger_words(versions) + return context + + def _base_model_aliases(data: dict[str, Any]) -> list[str]: """Return the site's own names for the base model. diff --git a/tests/services/test_agent_enrichment_source.py b/tests/services/test_agent_enrichment_source.py index db8a9a93..5d69bf12 100644 --- a/tests/services/test_agent_enrichment_source.py +++ b/tests/services/test_agent_enrichment_source.py @@ -104,8 +104,10 @@ class TestBuildPromptContext: mock_fetch.assert_awaited_once_with("jj3550945163/Krea-2-LORA") # The per-file lookup must receive the basename, not the full path. - mock_context.assert_awaited_once_with( - "jj3550945163/Krea-2-LORA", "krea.safetensors" + mock_context.assert_awaited_once() + assert mock_context.call_args.args == ( + "jj3550945163/Krea-2-LORA", + "krea.safetensors", ) assert context["source_platform"] == "modelscope" assert context["source_id"] == "jj3550945163/Krea-2-LORA" @@ -444,7 +446,8 @@ class TestLoadSourceCard: }, ) - mock_ctx.assert_awaited_once_with("u/r", "krea.safetensors") + mock_ctx.assert_awaited_once() + assert mock_ctx.call_args.args == ("u/r", "krea.safetensors") assert context is card assert variables["readme_content_full"] == "# card" assert variables["source_description"] == "作者说明" diff --git a/tests/services/test_model_sources.py b/tests/services/test_model_sources.py index 37cd7f9e..db59370f 100644 --- a/tests/services/test_model_sources.py +++ b/tests/services/test_model_sources.py @@ -10,6 +10,7 @@ from __future__ import annotations import pytest from py.services.model_sources import ( + ModelSourceCache, HuggingFaceSource, ModelScopeSource, TensorArtSource, @@ -688,3 +689,113 @@ class TestDownloadSourceRegistry: assert get_download_source("nope") is None assert get_download_source("modelscope").platform == "modelscope" assert get_download_source("huggingface").platform == "huggingface" + + +# --------------------------------------------------------------------------- +# Per-run model-card cache +# --------------------------------------------------------------------------- + + +class TestModelSourceCache: + def test_starts_empty(self): + cache = ModelSourceCache() + assert cache.readmes == {} + assert cache.provider == {} + + +class TestCachedModelCardFetch: + @pytest.mark.asyncio + async def test_detail_payload_is_fetched_once_per_source_id(self, monkeypatch): + """A collection repo's files share one detail request, not one each.""" + calls: list[str] = [] + + async def fake_fetch_json(url, **_kwargs): + calls.append(url) + return 200, _modelscope_detail_payload() + + monkeypatch.setattr( + "py.services.model_sources.modelscope.fetch_json", fake_fetch_json + ) + + source = ModelScopeSource() + cache = ModelSourceCache() + filenames = [ + "Krea-2-LORA_c1-st1000.safetensors", + "Krea-2-LORA_c1-st8000.safetensors", + "Krea-2-LORA_c1-st1000.safetensors", + ] + for name in filenames: + await source.fetch_model_card_context("u/r", name, cache=cache) + + assert calls == ["https://modelscope.cn/api/v1/models/u/r"] + assert ("modelscope", "detail", "u/r") in cache.provider + + @pytest.mark.asyncio + async def test_per_file_selection_still_runs_for_each_file(self, monkeypatch): + """The cached payload must not leak one file's images to another.""" + + async def fake_fetch_json(url, **_kwargs): + return 200, _modelscope_detail_payload() + + monkeypatch.setattr( + "py.services.model_sources.modelscope.fetch_json", fake_fetch_json + ) + + source = ModelScopeSource() + cache = ModelSourceCache() + + st1000 = await source.fetch_model_card_context( + "u/r", "Krea-2-LORA_c1-st1000.safetensors", cache=cache + ) + st8000 = await source.fetch_model_card_context( + "u/r", "Krea-2-LORA_c1-st8000.safetensors", cache=cache + ) + + assert st1000.example_images == [ + "https://resources.modelscope.cn/cover-images/b.png", + "https://resources.modelscope.cn/cover-images/c.png", + ] + assert st8000.example_images == [ + "https://resources.modelscope.cn/cover-images/a.png" + ] + assert st8000.trigger_words == [] + + @pytest.mark.asyncio + async def test_failures_are_not_cached(self, monkeypatch): + """A transient error must be retried for the next file.""" + calls: list[str] = [] + + async def fake_fetch_json(url, **_kwargs): + calls.append(url) + return 500, None + + monkeypatch.setattr( + "py.services.model_sources.modelscope.fetch_json", fake_fetch_json + ) + + source = ModelScopeSource() + cache = ModelSourceCache() + for _ in range(2): + context = await source.fetch_model_card_context("u/r", "a.safetensors", cache=cache) + assert context.is_empty() + + assert len(calls) == 2 + assert cache.provider == {} + + @pytest.mark.asyncio + async def test_no_cache_keeps_the_uncached_behaviour(self, monkeypatch): + calls: list[str] = [] + + async def fake_fetch_json(url, **_kwargs): + calls.append(url) + return 200, _modelscope_detail_payload() + + monkeypatch.setattr( + "py.services.model_sources.modelscope.fetch_json", fake_fetch_json + ) + + source = ModelScopeSource() + for _ in range(2): + await source.fetch_model_card_context("u/r", "a.safetensors") + + assert len(calls) == 2