Do not cache transient repository-list failures as empty weight sets. Deduplicate normalized file paths without lowering reported file sizes. Record whitelisted OOM evidence and bounded discovery diagnostics.
197 lines
7.8 KiB
Python
197 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
|
if str(PACKAGE_DIR) in sys.path:
|
|
sys.path.remove(str(PACKAGE_DIR))
|
|
sys.path.insert(0, str(PACKAGE_DIR))
|
|
|
|
from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree # noqa: E402
|
|
from http_json import HttpJsonError # noqa: E402
|
|
|
|
|
|
class RecordingHttpClient:
|
|
def __init__(self) -> None:
|
|
self.queries: list[dict] = []
|
|
|
|
def request_json(self, _method: str, _path: str, *, query: dict) -> dict:
|
|
self.queries.append(query)
|
|
return {"success": True, "data": {"models": []}}
|
|
|
|
|
|
class RateLimitedPageClient:
|
|
def __init__(self) -> None:
|
|
self.page_calls: list[int] = []
|
|
self.page_two_attempts = 0
|
|
|
|
def request_json(self, _method: str, _path: str, *, query: dict) -> dict:
|
|
page = int(query["page_number"])
|
|
self.page_calls.append(page)
|
|
if page == 2:
|
|
self.page_two_attempts += 1
|
|
if self.page_two_attempts == 1:
|
|
raise HttpJsonError("rate limited", status_code=429)
|
|
start = (page - 1) * 50
|
|
items = [
|
|
{
|
|
"id": f"owner/model-{index}",
|
|
"downloads": 100,
|
|
"last_modified": "2026-01-01T00:00:00Z",
|
|
}
|
|
for index in range(start, start + 50)
|
|
]
|
|
return {"success": True, "data": {"models": items}}
|
|
|
|
|
|
class ModelMetadataClient:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def request_json(self, _method: str, _path: str, **_kwargs) -> dict: # noqa: ANN003
|
|
self.calls += 1
|
|
return {"Code": 200, "Data": {"LastUpdatedTime": 1786294389}}
|
|
|
|
|
|
class ModelScopeDiscoveryTests(unittest.TestCase):
|
|
def test_listing_diagnostics_distinguish_filtered_prefix_from_exhaustion(self) -> None:
|
|
class FilteredPageClient:
|
|
def request_json(self, *_args, **_kwargs):
|
|
return {"data": {"models": [
|
|
{"id": f"owner/model-{index}", "downloads": 0}
|
|
for index in range(50)
|
|
]}}
|
|
|
|
discovery = HuggingFaceDiscovery(http_client=FilteredPageClient()) # type: ignore[arg-type]
|
|
self.assertEqual([], discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"], limit=50, min_downloads=50,
|
|
))
|
|
diagnostic = discovery.last_listing_diagnostics[0]
|
|
self.assertEqual(50, diagnostic["rawRowsScanned"])
|
|
self.assertEqual(0, diagnostic["acceptedModels"])
|
|
self.assertEqual("page_budget", diagnostic["stopReason"])
|
|
self.assertFalse(diagnostic["catalogExhausted"])
|
|
|
|
def test_failed_repo_tree_is_retried_instead_of_cached_as_empty(self) -> None:
|
|
class RecoveringClient:
|
|
calls = 0
|
|
|
|
def request_json(self, *_args, **_kwargs):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
raise HttpJsonError("temporary outage", status_code=503)
|
|
return {"Data": {"Files": [{"Path": "model.safetensors", "Size": 1024}]}}
|
|
|
|
client = RecoveringClient()
|
|
discovery = HuggingFaceDiscovery(legacy_http_client=client) # type: ignore[arg-type]
|
|
with self.assertRaises(HttpJsonError):
|
|
discovery.list_repo_tree("owner/recoverable")
|
|
self.assertNotIn("owner/recoverable", discovery._repo_tree_cache)
|
|
entries = discovery.list_repo_tree("owner/recoverable")
|
|
self.assertEqual("model.safetensors", entries[0]["Path"])
|
|
self.assertEqual(entries, discovery.list_repo_tree("owner/recoverable"))
|
|
self.assertEqual(2, client.calls)
|
|
|
|
def test_duplicate_normalized_paths_do_not_inflate_repository_size(self) -> None:
|
|
gib = 1024 ** 3
|
|
inspection = inspect_repo_tree("owner/model", [
|
|
{"Path": "model.safetensors", "Size": 20 * gib},
|
|
{"Path": "./model.safetensors", "Size": 19 * gib},
|
|
{"Path": "/model.safetensors", "Size": 20 * gib},
|
|
])
|
|
self.assertEqual(["model.safetensors"], inspection.file_paths)
|
|
self.assertEqual(["model.safetensors"], inspection.weight_files)
|
|
self.assertEqual(20 * gib, inspection.repository_size_bytes)
|
|
self.assertEqual(20 * gib, inspection.estimated_load_bytes("vllm"))
|
|
|
|
def test_model_last_modified_uses_detail_api_and_is_cached(self) -> None:
|
|
metadata_client = ModelMetadataClient()
|
|
discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type]
|
|
|
|
first = discovery.get_model_last_modified("owner/model")
|
|
second = discovery.get_model_last_modified("owner/model")
|
|
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(1786294389, int(first.timestamp())) # type: ignore[union-attr]
|
|
self.assertEqual(1, metadata_client.calls)
|
|
|
|
def test_model_detail_cache_evicts_old_entries_at_fixed_limit(self) -> None:
|
|
metadata_client = ModelMetadataClient()
|
|
with patch.dict("os.environ", {"MODELSCOPE_DETAIL_CACHE_MAX_MODELS": "32"}):
|
|
discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type]
|
|
for index in range(40):
|
|
discovery.get_model_last_modified(f"owner/model-{index}")
|
|
|
|
self.assertEqual(32, len(discovery._model_last_modified_cache))
|
|
discovery.get_model_last_modified("owner/model-0")
|
|
self.assertEqual(41, metadata_client.calls)
|
|
|
|
def test_openapi_page_size_never_exceeds_platform_limit(self) -> None:
|
|
http_client = RecordingHttpClient()
|
|
discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type]
|
|
|
|
models = discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"],
|
|
limit=1000,
|
|
min_downloads=0,
|
|
)
|
|
|
|
self.assertEqual([], models)
|
|
self.assertEqual(50, http_client.queries[0]["page_size"])
|
|
|
|
def test_popular_sort_uses_separate_page_cache(self) -> None:
|
|
http_client = RecordingHttpClient()
|
|
discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type]
|
|
|
|
discovery.list_recent_models(pipeline_tags=["text-generation"], limit=50, min_downloads=0)
|
|
discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"], limit=50, min_downloads=0, sort="downloads"
|
|
)
|
|
self.assertEqual(["last_modified", "downloads"], [item["sort"] for item in http_client.queries])
|
|
|
|
def test_popular_scan_rotates_bounded_page_windows(self) -> None:
|
|
http_client = RecordingHttpClient()
|
|
discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type]
|
|
|
|
for _ in range(2):
|
|
discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"],
|
|
limit=200,
|
|
min_downloads=50,
|
|
sort="downloads",
|
|
rotate_pages=True,
|
|
)
|
|
|
|
self.assertEqual([1, 5], [item["page_number"] for item in http_client.queries])
|
|
|
|
def test_rate_limited_page_keeps_previous_pages_and_resumes_from_cache(self) -> None:
|
|
http_client = RateLimitedPageClient()
|
|
discovery = HuggingFaceDiscovery(
|
|
http_client=http_client, # type: ignore[arg-type]
|
|
page_interval_seconds=0,
|
|
page_cache_ttl_seconds=60,
|
|
)
|
|
|
|
partial = discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"],
|
|
limit=100,
|
|
min_downloads=0,
|
|
)
|
|
resumed = discovery.list_recent_models(
|
|
pipeline_tags=["text-generation"],
|
|
limit=100,
|
|
min_downloads=0,
|
|
)
|
|
|
|
self.assertEqual(50, len(partial))
|
|
self.assertEqual(100, len(resumed))
|
|
self.assertEqual([1, 2, 2], http_client.page_calls)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|