adapt for ModelHub agent platform
This commit is contained in:
CoolBoy
169
modelhub_submmit_api/template_selector.py
Normal file
169
modelhub_submmit_api/template_selector.py
Normal file
@@ -0,0 +1,169 @@
|
||||
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
|
||||
Reference in New Issue
Block a user