110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from app.clients.modelhub import ModelHubClient
|
|
from app.domain.machines import machine_to_gpu_alias, parse_machine_names
|
|
from app.domain.model_type import engine_for_model
|
|
from app.domain.priorities import PRIORITY_BOUNTY
|
|
from app.settings import Settings
|
|
from app.storage.repositories import Repository
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
def _model_task_name(model: dict) -> str:
|
|
task_info = model.get("taskLevelsInfo") or {}
|
|
return model.get("taskLevelChineseName") or task_info.get("taskLevelChineseName") or ""
|
|
|
|
|
|
def _adapted_machine_names(model: dict) -> set[str]:
|
|
machines = model.get("machines") or []
|
|
return {m.get("machineName") for m in machines if m.get("machineName")}
|
|
|
|
|
|
class NotAdaptedCrawler:
|
|
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
|
|
self.settings = settings
|
|
self.repo = repo
|
|
self.client = client
|
|
|
|
def refresh(self) -> int:
|
|
if not self.settings.enable_not_adapted_crawler:
|
|
return 0
|
|
machine_names = parse_machine_names(self.settings.target_machine_names)
|
|
if not machine_names:
|
|
LOG.info("not-adapted crawler skipped: TARGET_MACHINE_NAMES is empty")
|
|
return 0
|
|
|
|
machine_alias_pairs: list[tuple[str, str]] = []
|
|
for machine_name in machine_names:
|
|
alias = machine_to_gpu_alias(machine_name)
|
|
if not alias:
|
|
LOG.warning("unknown target machine name, skipped: %s", machine_name)
|
|
continue
|
|
machine_alias_pairs.append((machine_name, alias))
|
|
if not machine_alias_pairs:
|
|
return 0
|
|
|
|
imported = 0
|
|
page = 1
|
|
while True:
|
|
data = self.client.list_modelhub_models(page, self.settings.modelhub_page_size, "")
|
|
records = data.get("records") or []
|
|
total = int(data.get("total") or 0)
|
|
if not records:
|
|
break
|
|
|
|
for model in records:
|
|
model_id = model.get("modelId")
|
|
if not model_id:
|
|
continue
|
|
task_name = _model_task_name(model)
|
|
if self.settings.target_task_level and self.settings.target_task_level not in task_name:
|
|
continue
|
|
|
|
adapted = _adapted_machine_names(model)
|
|
engine = engine_for_model(model_id)
|
|
for machine_name, alias in machine_alias_pairs:
|
|
if machine_name in adapted:
|
|
continue
|
|
if engine != "vllm":
|
|
self.repo.upsert_candidate(
|
|
model_id=model_id,
|
|
source="HUGGING_FACE",
|
|
origin="bounty_target_not_adapted_gguf",
|
|
priority=PRIORITY_BOUNTY,
|
|
is_bounty=True,
|
|
is_promote=True,
|
|
known_downloaded_by_others=True,
|
|
self_download_allowed=False,
|
|
target_gpu_aliases=alias,
|
|
notes=f"GGUF/llamacpp candidate for {machine_name}; vLLM phase skips submit",
|
|
)
|
|
imported += 1
|
|
continue
|
|
self.repo.upsert_candidate(
|
|
model_id=model_id,
|
|
source="HUGGING_FACE",
|
|
origin="bounty_target_not_adapted",
|
|
priority=PRIORITY_BOUNTY,
|
|
is_bounty=True,
|
|
is_promote=True,
|
|
known_downloaded_by_others=True,
|
|
self_download_allowed=False,
|
|
target_gpu_aliases=alias,
|
|
notes=f"not adapted on {machine_name}",
|
|
)
|
|
imported += 1
|
|
|
|
if self.settings.crawler_max_pages and page >= self.settings.crawler_max_pages:
|
|
break
|
|
if total and page * self.settings.modelhub_page_size >= total:
|
|
break
|
|
if len(records) < self.settings.modelhub_page_size:
|
|
break
|
|
page += 1
|
|
|
|
LOG.info("not-adapted crawler imported/updated %s candidates", imported)
|
|
return imported
|