Rebuild agent on original pooled runner

modelhub_submmit_api/hf_discovery.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
 import threading
 from pathlib import PurePosixPath
 from typing import Any
+from urllib.parse import quote
 
 from common import parse_datetime
+from defaults import EMBEDDED_MODELSCOPE_TOKEN
+from http_json import HttpJsonError
 from http_json import JsonHttpClient
 from models import HFModelSummary, ModelInspection
 
@@ -15,20 +17,47 @@ GGUF_PRIORITY = ("q4_0.gguf", "q8_0.gguf", "fp16.gguf")
 VLLM_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth")
 ONNX_WEIGHT_SUFFIXES = (".onnx",)
 
+MODELSCOPE_TASK_TAGS = {
+    "text-generation": "text-generation",
+    "image-text-to-text": "image-text-to-text",
+    "visual-question-answering": "visual-question-answering",
+    "document-question-answering": "document-question-answering",
+    "video-text-to-text": "video-text-to-text",
+    "text-to-image": "text-to-image-synthesis",
+    "image-to-image": "image-to-image",
+    "automatic-speech-recognition": "auto-speech-recognition",
+    "question-answering": "question-answering",
+    "feature-extraction": "feature-extraction",
+    "sentence-similarity": "sentence-similarity",
+    "image-classification": "image-classification",
+    "zero-shot-image-classification": "zero-shot-image-classification",
+    "text-classification": "text-classification",
+    "zero-shot-classification": "zero-shot-classification",
+    "reinforcement-learning": "reinforcement-learning",
+}
+
 
 class HuggingFaceDiscovery:
     def __init__(
         self,
-        base_url: str = "https://huggingface.co",
+        base_url: str = "https://modelscope.cn",
         http_client: JsonHttpClient | None = None,
+        legacy_http_client: JsonHttpClient | None = None,
         timeout: int = 30,
         retries: int = 2,
     ) -> None:
-        hf_token = os.getenv("HF_TOKEN")
+        token = os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN
         headers = {"User-Agent": "modelhub-submmit-cli/0.1"}
-        if hf_token:
-            headers["Authorization"] = f"Bearer {hf_token}"
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+            headers["Cookie"] = f"m_session_id={token}"
         self.http_client = http_client or JsonHttpClient(
+            base_url=f"{base_url.rstrip('/')}/openapi/v1",
+            default_headers=headers,
+            timeout=timeout,
+            retries=retries,
+        )
+        self.legacy_http_client = legacy_http_client or JsonHttpClient(
             base_url=base_url,
             default_headers=headers,
             timeout=timeout,
@@ -50,36 +79,15 @@ class HuggingFaceDiscovery:
             return []
 
         deduped: dict[str, HFModelSummary] = {}
-        max_workers = min(len(pipeline_tags), max(1, read_concurrency))
+        del read_concurrency
 
-        if max_workers <= 1:
-            tag_results = [
-                self._query_recent_models(
-                    pipeline_tag=pipeline_tag,
-                    limit=limit,
-                    min_downloads=min_downloads,
-                    updated_after=updated_after,
-                )
-                for pipeline_tag in pipeline_tags
-            ]
-        else:
-            with ThreadPoolExecutor(max_workers=max_workers) as executor:
-                futures = {
-                    executor.submit(
-                        self._query_recent_models,
-                        pipeline_tag=pipeline_tag,
-                        limit=limit,
-                        min_downloads=min_downloads,
-                        updated_after=updated_after,
-                    ): pipeline_tag
-                    for pipeline_tag in pipeline_tags
-                }
-                tag_results = []
-                for future in as_completed(futures):
-                    tag_results.append(future.result())
-
-        for models in tag_results:
-            for model in models:
+        for pipeline_tag in pipeline_tags:
+            for model in self._query_recent_models(
+                pipeline_tag=pipeline_tag,
+                limit=limit,
+                min_downloads=min_downloads,
+                updated_after=updated_after,
+            ):
                 current = deduped.get(model.repo_id)
                 if current is None or (model.last_modified or parse_datetime("1970-01-01")) > (
                     current.last_modified or parse_datetime("1970-01-01")
@@ -97,23 +105,42 @@ class HuggingFaceDiscovery:
         min_downloads: int,
         updated_after=None,
     ) -> list[HFModelSummary]:
-        payload = self.http_client.request_json(
-            "GET",
-            "/api/models",
-            query={
-                "pipeline_tag": pipeline_tag,
-                "sort": "lastModified",
-                "direction": "-1",
-                "limit": limit,
-                "full": "true",
-            },
-        )
-
+        page_size = min(max(1, limit), 100)
+        max_items = min(max(1, limit), 3000)
+        task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag)
         models: list[HFModelSummary] = []
-        for item in payload or []:
-            model = self._parse_model(item, min_downloads=min_downloads, updated_after=updated_after)
-            if model is not None:
-                models.append(model)
+        for page_number in range(1, (max_items + page_size - 1) // page_size + 1):
+            try:
+                payload = self.http_client.request_json(
+                    "GET",
+                    "/models",
+                    query={
+                        "page_number": page_number,
+                        "page_size": page_size,
+                        "sort": "last_modified",
+                        "filter.task": task_tag,
+                    },
+                )
+            except HttpJsonError as exc:
+                print(f"[modelscope] list_models_error task={task_tag} page={page_number} error={exc}", flush=True)
+                break
+
+            items = self._extract_models(payload)
+            if not items:
+                break
+            for item in items:
+                model = self._parse_model(
+                    item,
+                    fallback_pipeline_tag=pipeline_tag,
+                    min_downloads=min_downloads,
+                    updated_after=updated_after,
+                )
+                if model is not None:
+                    models.append(model)
+                if len(models) >= max_items:
+                    return models
+            if len(items) < page_size:
+                break
         return models
 
     def list_recent_text_generation_models(
@@ -140,40 +167,67 @@ class HuggingFaceDiscovery:
         if cached is not None:
             return list(cached)
 
-        payload = self.http_client.request_json(
-            "GET",
-            f"/api/models/{repo_id}/tree/main",
-            query={"recursive": "1"},
-        )
+        encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
+        try:
+            payload = self.legacy_http_client.request_json(
+                "GET",
+                f"/api/v1/models/{encoded_repo_id}/repo/files",
+                query={"Revision": "master", "Recursive": "true"},
+            )
+        except HttpJsonError as exc:
+            print(f"[modelscope] repo_tree_error repo={repo_id} error={exc}", flush=True)
+            payload = None
 
-        entries: list[dict[str, Any]]
-        if isinstance(payload, list):
-            entries = payload
-        elif isinstance(payload, dict):
-            for key in ("items", "tree", "siblings"):
-                value = payload.get(key)
-                if isinstance(value, list):
-                    entries = value
-                    break
-            else:
-                entries = []
-        else:
-            entries = []
+        entries = self._extract_files(payload)
 
         with self._repo_tree_lock:
             self._repo_tree_cache[repo_id] = list(entries)
         return list(entries)
 
     @staticmethod
-    def _parse_model(item: dict[str, Any], *, min_downloads: int, updated_after=None) -> HFModelSummary | None:
-        repo_id = item.get("id") or item.get("modelId")
+    def _extract_models(payload: Any) -> list[dict[str, Any]]:
+        data = payload.get("data") if isinstance(payload, dict) else payload
+        if isinstance(data, dict):
+            for key in ("models", "Models", "items", "list", "data", "results"):
+                value = data.get(key)
+                if isinstance(value, list):
+                    return [item for item in value if isinstance(item, dict)]
+        if isinstance(data, list):
+            return [item for item in data if isinstance(item, dict)]
+        return []
+
+    @staticmethod
+    def _extract_files(payload: Any) -> list[dict[str, Any]]:
+        data = payload.get("Data") if isinstance(payload, dict) else payload
+        if isinstance(data, dict):
+            for key in ("Files", "files", "items", "tree"):
+                value = data.get(key)
+                if isinstance(value, list):
+                    return [item for item in value if isinstance(item, dict)]
+        if isinstance(data, list):
+            return [item for item in data if isinstance(item, dict)]
+        return []
+
+    @staticmethod
+    def _parse_model(
+        item: dict[str, Any],
+        *,
+        fallback_pipeline_tag: str,
+        min_downloads: int,
+        updated_after=None,
+    ) -> HFModelSummary | None:
+        repo_id = item.get("id") or item.get("model_id") or item.get("modelId")
+        if not repo_id:
+            owner = item.get("owner") or item.get("Owner") or item.get("Path")
+            name = item.get("name") or item.get("Name")
+            repo_id = f"{owner}/{name}" if owner and name else None
         if not repo_id:
             return None
-        downloads = int(item.get("downloads") or 0)
+        downloads = int(item.get("downloads") or item.get("Downloads") or 0)
         if downloads < min_downloads:
             return None
-        pipeline_tag = item.get("pipeline_tag") or item.get("pipelineTag")
-        last_modified = parse_datetime(item.get("lastModified") or item.get("last_modified"))
+        pipeline_tag = fallback_pipeline_tag
+        last_modified = parse_datetime(item.get("last_modified") or item.get("UpdatedAt") or item.get("LastUpdatedTime"))
         if updated_after and last_modified and last_modified < updated_after:
             return None
         return HFModelSummary(
@@ -181,7 +235,7 @@ class HuggingFaceDiscovery:
             downloads=downloads,
             last_modified=last_modified,
             pipeline_tag=pipeline_tag,
-            created_at=parse_datetime(item.get("createdAt") or item.get("created_at")),
+            created_at=parse_datetime(item.get("created_at") or item.get("CreatedAt")),
         )