feat: add tiered durable state and memory bounds

modelhub_submmit_api/hf_discovery.py

@@ -4,6 +4,7 @@ import os
 import re
 import threading
 import time
+from collections import OrderedDict
 from dataclasses import replace
 from datetime import datetime, timezone
 from pathlib import PurePosixPath
@@ -70,13 +71,23 @@ class HuggingFaceDiscovery:
             timeout=timeout,
             retries=retries,
         )
-        self._repo_tree_cache: dict[str, list[dict[str, Any]]] = {}
+        self._detail_cache_limit = max(
+            32,
+            min(1000, int(os.getenv("MODELSCOPE_DETAIL_CACHE_MAX_MODELS", "128"))),
+        )
+        self._page_cache_limit = max(
+            16,
+            min(500, int(os.getenv("MODELSCOPE_PAGE_CACHE_MAX_PAGES", "96"))),
+        )
+        self._repo_tree_cache: OrderedDict[str, list[dict[str, Any]]] = OrderedDict()
         self._repo_tree_lock = threading.Lock()
-        self._model_config_cache: dict[str, tuple[dict[str, Any], str | None]] = {}
+        self._model_config_cache: OrderedDict[str, tuple[dict[str, Any], str | None]] = OrderedDict()
         self._model_config_lock = threading.Lock()
-        self._model_last_modified_cache: dict[str, datetime | None] = {}
+        self._model_last_modified_cache: OrderedDict[str, datetime | None] = OrderedDict()
         self._model_last_modified_lock = threading.Lock()
-        self._model_page_cache: dict[tuple[str, int, int], tuple[float, list[dict[str, Any]]]] = {}
+        self._model_page_cache: OrderedDict[
+            tuple[str, int, int], tuple[float, list[dict[str, Any]]]
+        ] = OrderedDict()
         self._model_page_cache_ttl = max(
             0.0,
             float(
@@ -95,9 +106,16 @@ class HuggingFaceDiscovery:
         )
         self._last_model_page_request_at = 0.0
         self._unsupported_task_filters: dict[str, float] = {}
-        self._model_card_cache: dict[str, dict[str, Any]] = {}
+        self._model_card_cache: OrderedDict[str, dict[str, Any]] = OrderedDict()
         self._model_card_lock = threading.Lock()
 
+    @staticmethod
+    def _remember_lru(cache: OrderedDict, key: Any, value: Any, limit: int) -> None:
+        cache[key] = value
+        cache.move_to_end(key)
+        while len(cache) > limit:
+            cache.popitem(last=False)
+
     def list_recent_models(
         self,
         *,
@@ -149,6 +167,7 @@ class HuggingFaceDiscovery:
             cache_key = ("*" if filter_disabled else task_tag, page_number, page_size)
             cached = self._model_page_cache.get(cache_key)
             if cached is not None and time.monotonic() - cached[0] < self._model_page_cache_ttl:
+                self._model_page_cache.move_to_end(cache_key)
                 items = list(cached[1])
             else:
                 elapsed = time.monotonic() - self._last_model_page_request_at
@@ -191,7 +210,12 @@ class HuggingFaceDiscovery:
                     break
 
                 items = self._extract_models(payload)
-                self._model_page_cache[cache_key] = (time.monotonic(), list(items))
+                self._remember_lru(
+                    self._model_page_cache,
+                    cache_key,
+                    (time.monotonic(), list(items)),
+                    self._page_cache_limit,
+                )
             if not items:
                 break
             for item in items:
@@ -247,6 +271,8 @@ class HuggingFaceDiscovery:
     def get_model_card_metadata(self, repo_id: str) -> dict[str, Any]:
         with self._model_card_lock:
             cached = self._model_card_cache.get(repo_id)
+            if cached is not None:
+                self._model_card_cache.move_to_end(repo_id)
         if cached is not None:
             return dict(cached)
         encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
@@ -262,12 +288,19 @@ class HuggingFaceDiscovery:
         except Exception:
             result = {}
         with self._model_card_lock:
-            self._model_card_cache[repo_id] = dict(result)
+            self._remember_lru(
+                self._model_card_cache,
+                repo_id,
+                dict(result),
+                self._detail_cache_limit,
+            )
         return result
 
     def get_model_config(self, repo_id: str) -> tuple[dict[str, Any], str | None]:
         with self._model_config_lock:
             cached = self._model_config_cache.get(repo_id)
+            if cached is not None:
+                self._model_config_cache.move_to_end(repo_id)
         if cached is not None:
             return dict(cached[0]), cached[1]
 
@@ -284,13 +317,20 @@ class HuggingFaceDiscovery:
             result = ({}, f"{type(exc).__name__}: {exc}")
 
         with self._model_config_lock:
-            self._model_config_cache[repo_id] = result
+            self._remember_lru(
+                self._model_config_cache,
+                repo_id,
+                result,
+                self._detail_cache_limit,
+            )
         return dict(result[0]), result[1]
 
     def get_model_last_modified(self, repo_id: str) -> datetime | None:
         with self._model_last_modified_lock:
             if repo_id in self._model_last_modified_cache:
-                return self._model_last_modified_cache[repo_id]
+                result = self._model_last_modified_cache[repo_id]
+                self._model_last_modified_cache.move_to_end(repo_id)
+                return result
 
         encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
         try:
@@ -325,12 +365,19 @@ class HuggingFaceDiscovery:
                 result = parse_datetime(raw_value)
 
         with self._model_last_modified_lock:
-            self._model_last_modified_cache[repo_id] = result
+            self._remember_lru(
+                self._model_last_modified_cache,
+                repo_id,
+                result,
+                self._detail_cache_limit,
+            )
         return result
 
     def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
         with self._repo_tree_lock:
             cached = self._repo_tree_cache.get(repo_id)
+            if cached is not None:
+                self._repo_tree_cache.move_to_end(repo_id)
         if cached is not None:
             return list(cached)
 
@@ -348,7 +395,12 @@ class HuggingFaceDiscovery:
         entries = self._extract_files(payload)
 
         with self._repo_tree_lock:
-            self._repo_tree_cache[repo_id] = list(entries)
+            self._remember_lru(
+                self._repo_tree_cache,
+                repo_id,
+                list(entries),
+                self._detail_cache_limit,
+            )
         return list(entries)
 
     @staticmethod