submmit/modelhub_submmit_api/template_selector.py

from __future__ import annotations

import json
import re
import os
from pathlib import Path

from models import TemplateRecord


MODULE_DIR = Path(__file__).resolve().parent
DEFAULT_TEMPLATE_FILE_NAME = "adapt_task_templates.jsonl"
DEFAULT_TEMPLATE_RELATIVE_CANDIDATES = (
    Path("templates/public_submit") / DEFAULT_TEMPLATE_FILE_NAME,
    Path("model adaptation/templates/public_submit") / DEFAULT_TEMPLATE_FILE_NAME,
)


def _normalize_token(value: str) -> str:
    return re.sub(r"[^a-z0-9]", "", value.lower())


GPU_ALIASES = {
    "k100": "hygon_k100-ai",
    "k100ai": "hygon_k100-ai",
    "hygonk100ai": "hygon_k100-ai",
    "hygon_k100ai": "hygon_k100-ai",
    "hygonk100": "hygon_k100-ai",
    "s4000": "Mthreads_s4000",
    "mthreadss4000": "Mthreads_s4000",
    "bi150": "Iluvatar_bi-150",
    "iluvatarbi150": "Iluvatar_bi-150",
    "bi100": "Iluvatar_bi-100",
    "iluvatarbi100": "Iluvatar_bi-100",
    "mrv100": "Iluvatar_mrv-100",
    "iluvatarmrv100": "Iluvatar_mrv-100",
    "metaxc500": "MetaX_c-500",
    "c500": "MetaX_c-500",
    "kunlunxinp800": "Kunlunxin_p-800",
    "p800": "Kunlunxin_p-800",
    "kunlunxinr2008f": "Kunlunxin_r-200-8f",
    "r2008f": "Kunlunxin_r-200-8f",
    "ascend910b3": "Ascend_910-b3",
    "910b3": "Ascend_910-b3",
    "ascend910b4": "Ascend_910-b4",
    "910b4": "Ascend_910-b4",
    "bieren166m": "Biren_166m",
    "biren166m": "Biren_166m",
    "vastaiva16": "Vastai_va16",
    "va16": "Vastai_va16",
    "sunrisept200x1": "Sunrise_pt-200-x1",
    "pt200x1": "Sunrise_pt-200-x1",
}


class TemplateSelector:
    def __init__(self, template_file: Path | None = None) -> None:
        self.template_file = self._resolve_template_file(template_file)
        self.templates = self._load_templates(self.template_file)
        self.index = {
            (template.task_type, template.framework, template.target_gpu): template
            for template in self.templates
        }
        for template in self.templates:
            GPU_ALIASES.setdefault(_normalize_token(template.target_gpu), template.target_gpu)

    @staticmethod
    def _resolve_template_file(template_file: Path | None) -> Path:
        env_template_file = os.getenv("MODELHUB_TEMPLATE_FILE")
        if env_template_file:
            candidate = Path(env_template_file).expanduser()
            if candidate.exists():
                return candidate

        candidates: list[Path] = []
        if template_file is not None:
            candidates.append(template_file)
        else:
            search_roots = [Path.cwd().resolve(), MODULE_DIR.resolve()]
            for root in search_roots:
                candidates.append(root / DEFAULT_TEMPLATE_FILE_NAME)
                candidates.extend(root / relative for relative in DEFAULT_TEMPLATE_RELATIVE_CANDIDATES)
                parent = root.parent
                if parent != root:
                    candidates.append(parent / DEFAULT_TEMPLATE_FILE_NAME)
                    candidates.extend(parent / relative for relative in DEFAULT_TEMPLATE_RELATIVE_CANDIDATES)

        for candidate in candidates:
            if candidate.exists():
                return candidate

        tried = ", ".join(str(candidate) for candidate in candidates)
        raise FileNotFoundError(f"Could not find ModelHub template file. Tried: {tried}")

    @staticmethod
    def _load_templates(path: Path) -> list[TemplateRecord]:
        templates: list[TemplateRecord] = []
        for line in path.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            payload = json.loads(line)
            source = payload.get("source") or {}
            templates.append(
                TemplateRecord(
                    template_id=payload["templateId"],
                    task_type=payload["taskType"],
                    framework=payload["framework"],
                    target_gpu=payload["targetGpu"],
                    model_address_template=payload["modelAddressTemplate"],
                    config_params=payload["configParams"],
                    source_file=source.get("file"),
                    source_line=source.get("line"),
                )
            )
        return templates

    def normalize_gpu(self, gpu: str) -> str:
        normalized = _normalize_token(gpu)
        if normalized in GPU_ALIASES:
            return GPU_ALIASES[normalized]
        raise KeyError(f"Unsupported GPU alias: {gpu}")

    def supported_frameworks(self, task_type: str, target_gpu: str) -> set[str]:
        return {
            template.framework
            for template in self.templates
            if template.task_type == task_type and template.target_gpu == target_gpu
        }

    def supported_target_gpus(self, task_type: str, *, auto_only: bool = False) -> list[str]:
        gpus = {
            template.target_gpu
            for template in self.templates
            if template.task_type == task_type and (not auto_only or self.is_auto_template(template))
        }
        return sorted(gpus)

    def supported_frameworks_for_auto(self, task_type: str, target_gpu: str) -> set[str]:
        return {
            template.framework
            for template in self.templates
            if template.task_type == task_type and template.target_gpu == target_gpu and self.is_auto_template(template)
        }

    def select_template(self, task_type: str, framework: str, target_gpu: str) -> TemplateRecord:
        key = (task_type, framework, target_gpu)
        if key not in self.index:
            raise KeyError(f"Missing template for taskType={task_type}, framework={framework}, targetGpu={target_gpu}")
        return self.index[key]

    @staticmethod
    def is_auto_template(template: TemplateRecord) -> bool:
        config = template.config_params
        if "PLACEHOLDER" in config:
            return False
        if "根据" in config:
            return False
        return True

    def render_config(self, template: TemplateRecord, gguf_filename: str | None = None) -> str:
        config = template.config_params
        if template.framework != "llamacpp":
            return config
        if not gguf_filename:
            raise ValueError("llamacpp template requires a GGUF filename")
        replaced = re.sub(r"/model/[^,\]\s]+?\.gguf", f"/model/{gguf_filename}", config)
        if replaced == config:
            raise ValueError(f"Failed to replace GGUF filename in template {template.template_id}")
        return replaced