diff --git a/src/modelscope_hub/_legacy_api.py b/src/modelscope_hub/_legacy_api.py index b460e33..2fc7547 100644 --- a/src/modelscope_hub/_legacy_api.py +++ b/src/modelscope_hub/_legacy_api.py @@ -15,6 +15,7 @@ from __future__ import annotations import uuid +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import IO, Any, BinaryIO from urllib.parse import quote_plus, urlparse @@ -26,6 +27,9 @@ API_MAX_RETRIES, API_TIMEOUT, LEGACY_API_PREFIX, + REPO_FILES_TRUNCATION_LIMIT, + REPO_TREE_MAX_REQUESTS, + REPO_TREE_WALK_WORKERS, UPLOAD_BLOB_CONNECT_TIMEOUT, UPLOAD_BLOB_READ_TIMEOUT, UPLOAD_RETRY_ALLOWED_METHODS, @@ -62,6 +66,21 @@ def _resolve_segment(repo_type: str) -> str: return _REPO_TYPE_SEGMENT.get(repo_type, f"{repo_type}s") +def _is_dataset(repo_type: str) -> bool: + """Whether the repo_type addresses the dataset endpoints.""" + return repo_type in (RepoType.DATASET, "dataset", "datasets") + + +def _entry_path(entry: dict) -> str: + """Return the repo-relative path of a file-tree entry.""" + return entry.get("Path") or entry.get("path") or entry.get("Name") or "" + + +def _is_dir_entry(entry: dict) -> bool: + """Whether a file-tree entry is a directory (git tree object).""" + return (entry.get("Type") or entry.get("type") or "blob") == "tree" + + class LegacyClient: """Internal client for /api/v1/ endpoints not covered by OpenAPI. @@ -316,7 +335,46 @@ def list_repo_files( Models/studios/etc: GET /api/v1/{type}s/{repo_id}/repo/files Datasets: GET /api/v1/datasets/{repo_id}/repo/tree + + ``repo/files`` supports no pagination and silently truncates at + :data:`REPO_FILES_TRUNCATION_LIMIT` entries, so a recursive listing that + comes back exactly at the limit is re-enumerated directory by directory + (see :meth:`_walk_repo_files`). Datasets take the ``repo/tree`` + endpoint, which does paginate, so they are paged through instead. """ + if _is_dataset(repo_type): + if recursive: + return self.list_dataset_files_paginated( + repo_id=repo_id, + revision=revision, + root_path=root or "/", + ) + return self._list_files_page(repo_id, repo_type, revision, recursive=False, root=root) + + entries = self._list_files_page(repo_id, repo_type, revision, recursive=recursive, root=root) + if len(entries) < REPO_FILES_TRUNCATION_LIMIT: + return entries + if not recursive: + logger.warning( + "Directory %r of %s holds at least %d entries and the server returns no more " + "than that for a single listing; the result is incomplete.", + root or "/", + repo_id, + REPO_FILES_TRUNCATION_LIMIT, + ) + return entries + return self._walk_repo_files(repo_id, repo_type, revision, root=root, prefetched=entries) + + def _list_files_page( + self, + repo_id: str, + repo_type: str, + revision: str, + *, + recursive: bool, + root: str | None = None, + ) -> list[dict]: + """Perform a single file-tree request and unwrap the entry list.""" segment = _resolve_segment(repo_type) params: dict[str, Any] = { "Revision": revision, @@ -325,7 +383,7 @@ def list_repo_files( if root: params["Root"] = root - suffix = "repo/tree" if repo_type in (RepoType.DATASET, "dataset", "datasets") else "repo/files" + suffix = "repo/tree" if _is_dataset(repo_type) else "repo/files" resp = self._request("GET", f"{segment}/{repo_id}/{suffix}", params=params) data = self._json_data(resp) if isinstance(data, list): @@ -335,6 +393,158 @@ def list_repo_files( return data.get("Files") or data.get("files") or [] return [] + def _walk_repo_files( + self, + repo_id: str, + repo_type: str, + revision: str, + *, + root: str | None = None, + prefetched: list[dict] | None = None, + ) -> list[dict]: + """Enumerate a file tree the server truncated, one directory at a time. + + Since ``repo/files`` caps every response, the only way to see past the + cap is to scope requests with ``Root``. Each subtree is first requested + whole; only the subtrees that come back at the cap are listed shallowly + and walked per child directory. The request count therefore scales with + the number of oversized directories, not with the directory total. + + The walk proceeds one tree level at a time and issues the listings of a + level concurrently, since a deep repo needs hundreds of round trips. + + ``prefetched`` is the already-truncated listing of ``root``, reused so + the caller's request is not repeated. Entries are de-duplicated by path. + Directories with more direct children than the cap cannot be enumerated + at all — those are reported through a warning and the returned list is + then knowingly incomplete. + """ + limit = REPO_FILES_TRUNCATION_LIMIT + collected: dict[str, dict] = {} + oversized: list[str] = [] + requests_made = 0 + progress_mark = 200 + budget_spent = False + + logger.info( + "Repo %s: file listing came back at the server cap of %d entries; walking the " + "tree per directory to recover the full list ...", + repo_id, + limit, + ) + + def absorb(entries: list[dict]) -> None: + for entry in entries: + path = _entry_path(entry) + if path: + collected.setdefault(path, entry) + + def fetch_level(roots: list[str | None], *, recursive: bool) -> dict[str | None, list[dict]]: + """List several directories at once, clamped to the request budget.""" + nonlocal requests_made, budget_spent, progress_mark + if not roots: + return {} + remaining = REPO_TREE_MAX_REQUESTS - requests_made + if remaining <= 0: + budget_spent = True + return {} + if len(roots) > remaining: + roots = roots[:remaining] + budget_spent = True + requests_made += len(roots) + + def one(subroot: str | None) -> list[dict]: + return self._list_files_page( + repo_id, + repo_type, + revision, + recursive=recursive, + root=subroot, + ) + + if len(roots) == 1: + listings = {roots[0]: one(roots[0])} + else: + workers = min(REPO_TREE_WALK_WORKERS, len(roots)) + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(one, subroot): subroot for subroot in roots} + listings = {futures[future]: future.result() for future in as_completed(futures)} + + if requests_made >= progress_mark: + logger.info( + "Repo %s: %d listings done, %d entries collected so far ...", + repo_id, + requests_made, + len(collected), + ) + progress_mark = requests_made + 200 + return listings + + frontier: list[str | None] = [root] + known: dict[str | None, list[dict]] = {} if prefetched is None else {root: prefetched} + + while frontier: + known.update(fetch_level([r for r in frontier if r not in known], recursive=True)) + + truncated: list[str | None] = [] + for subroot in frontier: + subtree = known.pop(subroot, None) + if subtree is None: + continue # budget ran out before this directory was reached + # A truncated subtree is still a valid prefix, so keep its entries + # and split the directory up to reach the rest. + absorb(subtree) + if len(subtree) >= limit: + truncated.append(subroot) + if not truncated or budget_spent: + break + + shallow_listings = fetch_level(truncated, recursive=False) + next_frontier: list[str | None] = [] + for subroot in truncated: + shallow = shallow_listings.get(subroot) + if shallow is None: + continue + absorb(shallow) + child_dirs = [path for path in map(_entry_path, filter(_is_dir_entry, shallow)) if path] + if len(shallow) >= limit or not child_dirs: + # This level alone exceeds the cap, or it has no sub-directories + # to split by: part of it stays invisible whatever we do. Still + # descend into the children we did see — that recovers strictly + # more than giving up here. + oversized.append(subroot or "/") + next_frontier.extend(child_dirs) + if budget_spent: + break + frontier = list(dict.fromkeys(next_frontier)) + + if budget_spent: + logger.warning( + "Repo %s: stopped walking the file tree after %d listings " + "(MODELSCOPE_REPO_TREE_MAX_REQUESTS=%d); the file list is incomplete.", + repo_id, + requests_made, + REPO_TREE_MAX_REQUESTS, + ) + if oversized: + shown = ", ".join(oversized[:5]) + (" ..." if len(oversized) > 5 else "") + logger.warning( + "Repo %s: %d director%s hold %d or more direct entries, which the server " + "cannot enumerate in full; the file list is incomplete: %s", + repo_id, + len(oversized), + "y" if len(oversized) == 1 else "ies", + limit, + shown, + ) + logger.info( + "Repo %s: collected %d entries from %d directory listings.", + repo_id, + len(collected), + requests_made, + ) + return list(collected.values()) + def list_dataset_files_paginated( self, repo_id: str, diff --git a/src/modelscope_hub/constants.py b/src/modelscope_hub/constants.py index aec9061..203d98b 100644 --- a/src/modelscope_hub/constants.py +++ b/src/modelscope_hub/constants.py @@ -295,6 +295,30 @@ def _env_register( "API_MAX_RETRIES", ) +REPO_FILES_TRUNCATION_LIMIT: int = 3000 +"""Server-side hard cap on a single ``repo/files`` listing. + +``GET /api/v1/{type}s/{repo_id}/repo/files`` silently truncates the file tree at +this many entries: the response is ``HTTP 200`` with ``Success: true``, carries +neither ``TotalCount`` nor a truncation flag, and ignores every pagination +parameter. A listing whose length equals this limit therefore means "there may +be more", and the tree has to be re-enumerated with ``Root``-scoped requests. +""" + +REPO_TREE_MAX_REQUESTS: int = _env_int( + "MODELSCOPE_REPO_TREE_MAX_REQUESTS", + 5000, + "Request budget for walking a truncated repo file tree", + "Network", +) + +REPO_TREE_WALK_WORKERS: int = _env_int( + "MODELSCOPE_REPO_TREE_WALK_WORKERS", + 8, + "Concurrent listings when walking a truncated repo file tree", + "Network", +) + # --------------------------------------------------------------------------- # Endpoint switching # --------------------------------------------------------------------------- diff --git a/tests/test_repo_files_truncation.py b/tests/test_repo_files_truncation.py new file mode 100644 index 0000000..088b1e5 --- /dev/null +++ b/tests/test_repo_files_truncation.py @@ -0,0 +1,243 @@ +"""Unit tests for the client-side mitigation of the ``repo/files`` entry cap. + +The server truncates every ``repo/files`` listing at +``REPO_FILES_TRUNCATION_LIMIT`` entries, returns ``HTTP 200`` with no truncation +marker, and ignores all pagination parameters. ``LegacyClient.list_repo_files`` +therefore treats a listing that lands exactly on the cap as "there may be more" +and re-enumerates the tree with ``Root``-scoped requests. + +Network-free: ``_list_files_page`` is replaced by an in-memory repo that mimics +the endpoint, cap included. The cap is patched down to a small number so the +fixtures stay readable. +""" + +from __future__ import annotations + +from unittest import mock + +import pytest + +from modelscope_hub import _legacy_api +from modelscope_hub._legacy_api import LegacyClient + +CAP = 5 + + +def _blob(path: str) -> dict: + return {"Path": path, "Name": path.rsplit("/", 1)[-1], "Type": "blob", "Size": 1} + + +def _tree(path: str) -> dict: + return {"Path": path, "Name": path.rsplit("/", 1)[-1], "Type": "tree", "Size": 0} + + +class _FakeRepo: + """Stand-in for ``repo/files``, including its silent truncation.""" + + def __init__(self, blobs: list[str], cap: int = CAP) -> None: + self.blobs = sorted(blobs) + self.cap = cap + self.calls: list[tuple[str | None, bool]] = [] + + def _entries(self, root: str | None, recursive: bool) -> list[dict]: + prefix = f"{root}/" if root else "" + found: dict[str, dict] = {} + for blob in self.blobs: + if not blob.startswith(prefix): + continue + rest = blob[len(prefix) :] + if "/" not in rest: + found[blob] = _blob(blob) + elif recursive: + parts = rest.split("/") + for depth in range(1, len(parts)): + nested = prefix + "/".join(parts[:depth]) + found[nested] = _tree(nested) + found[blob] = _blob(blob) + else: + child = prefix + rest.split("/")[0] + found[child] = _tree(child) + return list(found.values()) + + def listing(self, repo_id, repo_type, revision, *, recursive, root=None) -> list[dict]: + """Signature-compatible replacement for ``_list_files_page``.""" + self.calls.append((root, recursive)) + return self._entries(root, recursive)[: self.cap] + + def full_paths(self, recursive: bool = True) -> set[str]: + """Every path the endpoint would expose if it did not truncate.""" + return {_legacy_api._entry_path(e) for e in self._entries(None, recursive)} + + +@pytest.fixture +def client() -> LegacyClient: + return LegacyClient(token=None, endpoint="https://example.com") + + +def _install(repo: _FakeRepo): + """Patch the transport and shrink the cap to the fake repo's cap.""" + return ( + mock.patch.object(LegacyClient, "_list_files_page", side_effect=repo.listing), + mock.patch.object(_legacy_api, "REPO_FILES_TRUNCATION_LIMIT", repo.cap), + ) + + +class TestFastPath: + def test_small_repo_takes_exactly_one_request(self, client): + repo = _FakeRepo(["config.json", "model.safetensors", "sub/extra.bin"]) + transport, cap = _install(repo) + with transport, cap: + out = client.list_repo_files("owner/name", "model") + + assert len(repo.calls) == 1 + assert repo.calls[0] == (None, True) + assert {_legacy_api._entry_path(e) for e in out} == { + "config.json", + "model.safetensors", + "sub", + "sub/extra.bin", + } + + def test_listing_just_below_cap_is_not_rewalked(self, client): + # cap - 1 entries: a complete tree that must not trigger the fallback. + repo = _FakeRepo([f"f{i}.bin" for i in range(CAP - 1)]) + transport, cap = _install(repo) + with transport, cap: + out = client.list_repo_files("owner/name", "model") + + assert len(repo.calls) == 1 + assert len(out) == CAP - 1 + + +class TestTruncatedTreeIsRewalked: + def _repo(self) -> _FakeRepo: + # Root recursive listing hits the cap, so the tree must be walked: + # 3 top-level blobs + two directories holding 4 blobs each. + return _FakeRepo( + [ + "README.md", + "config.json", + "model.safetensors", + *[f"weights/part{i}.bin" for i in range(4)], + *[f"tokenizer/vocab{i}.txt" for i in range(4)], + ] + ) + + def test_full_tree_is_recovered(self, client): + repo = self._repo() + transport, cap = _install(repo) + with transport, cap: + out = client.list_repo_files("owner/name", "model") + + got = {_legacy_api._entry_path(e) for e in out} + assert got == repo.full_paths() + # The truncated single call would have returned only `cap` entries. + assert len(got) > repo.cap + + def test_walk_scopes_requests_with_root(self, client): + repo = self._repo() + transport, cap = _install(repo) + with transport, cap: + client.list_repo_files("owner/name", "model") + + roots_visited = {root for root, _ in repo.calls} + assert {"weights", "tokenizer"} <= roots_visited + # The caller's initial listing is reused, never repeated. + assert [call for call in repo.calls if call == (None, True)] == [(None, True)] + # The root level is listed shallowly to discover the child directories. + assert (None, False) in repo.calls + + def test_entries_are_deduplicated(self, client): + repo = self._repo() + transport, cap = _install(repo) + with transport, cap: + out = client.list_repo_files("owner/name", "model") + + paths = [_legacy_api._entry_path(e) for e in out] + assert len(paths) == len(set(paths)) + + def test_nested_oversized_subtree_is_split_further(self, client): + # `deep` alone exceeds the cap and only resolves via its children. + repo = _FakeRepo( + [ + "README.md", + *[f"deep/a/f{i}.bin" for i in range(4)], + *[f"deep/b/f{i}.bin" for i in range(4)], + ] + ) + transport, cap = _install(repo) + with transport, cap: + out = client.list_repo_files("owner/name", "model") + + assert {_legacy_api._entry_path(e) for e in out} == repo.full_paths() + assert {"deep", "deep/a", "deep/b"} <= {root for root, _ in repo.calls} + + +class TestIncompleteResultsAreReported: + @staticmethod + def _warnings(warn_mock) -> list[str]: + return [call.args[0] for call in warn_mock.call_args_list] + + def test_flat_oversized_directory_warns(self, client): + # A single directory with more direct blob children than the cap cannot + # be enumerated: no sub-directories exist to scope requests by. + repo = _FakeRepo([f"flat/f{i}.bin" for i in range(CAP + 3)]) + transport, cap = _install(repo) + with transport, cap, mock.patch.object(_legacy_api.logger, "warning") as warn: + out = client.list_repo_files("owner/name", "model") + + assert len(out) <= repo.cap + 1 # partial, best effort + assert any("incomplete" in message for message in self._warnings(warn)) + + def test_non_recursive_listing_at_cap_warns_without_walking(self, client): + repo = _FakeRepo([f"f{i}.bin" for i in range(CAP + 3)]) + transport, cap = _install(repo) + with transport, cap, mock.patch.object(_legacy_api.logger, "warning") as warn: + out = client.list_repo_files("owner/name", "model", recursive=False) + + assert len(repo.calls) == 1 # no fallback walk for a shallow listing + assert len(out) == repo.cap + assert any("incomplete" in message for message in self._warnings(warn)) + + def test_request_budget_stops_the_walk(self, client): + # Root holds 3 directories (a shallow listing fits under the cap), but the + # recursive listing does not — so the walk starts and then runs out of budget. + repo = _FakeRepo([f"d{d}/f{i}.bin" for d in range(3) for i in range(4)]) + transport, cap = _install(repo) + with ( + transport, + cap, + mock.patch.object(_legacy_api, "REPO_TREE_MAX_REQUESTS", 3), + mock.patch.object(_legacy_api.logger, "warning") as warn, + ): + out = client.list_repo_files("owner/name", "model") + + assert len(repo.calls) <= 1 + 3 # initial detection listing + the walk's budget + assert len(out) > 0 # whatever was collected is still returned + assert any("MODELSCOPE_REPO_TREE_MAX_REQUESTS" in message for message in self._warnings(warn)) + + +class TestDatasetRouting: + def test_recursive_dataset_listing_uses_the_paginated_endpoint(self, client): + pages = [{"Path": "data/train.csv", "Type": "blob", "Size": 1}] + with ( + mock.patch.object(LegacyClient, "list_dataset_files_paginated", return_value=pages) as paged, + mock.patch.object(LegacyClient, "_list_files_page") as single, + ): + out = client.list_repo_files("owner/ds", "dataset", revision="v1", root="data") + + assert out == pages + single.assert_not_called() + _, kwargs = paged.call_args + assert kwargs["revision"] == "v1" + assert kwargs["root_path"] == "data" + + def test_non_recursive_dataset_listing_stays_single_page(self, client): + with ( + mock.patch.object(LegacyClient, "list_dataset_files_paginated") as paged, + mock.patch.object(LegacyClient, "_list_files_page", return_value=[]) as single, + ): + client.list_repo_files("owner/ds", "dataset", recursive=False) + + paged.assert_not_called() + single.assert_called_once()