Bootstrap architecture rules from task history

This commit is contained in:
CoolBoy
2026-08-12 08:41:07 +08:00
parent 7ec875563e
commit d8ed7fa1c2
12 changed files with 653 additions and 63 deletions

View File

@@ -13,40 +13,10 @@ from modelhub_client import ModelHubClient, ModelHubClientPool
from models import HFModelSummary
from poll_runner import DEFAULT_POLL_RUNS_DIR, run_poll_loop
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
from task_registry import TASK_SPEC_BY_TYPE
from task_registry import TASK_SPEC_BY_TYPE, task_type_from_history_task
from template_selector import TemplateSelector
TASK_TYPE_BY_MODEL_TASK_LEVEL_ID = {
"0": "text-generation",
"2": "visual-multi-modal",
"21": "text-to-image-generation",
"23": "text-generation",
"29": "feature_emb",
"34": "feature_emb",
"39": "question_answering",
"54": "asr",
"116": "vision_classification",
"198": "vision_classification",
"207": "reinforcement_learning",
"213": "text-to-image-generation",
}
TASK_TYPE_BY_MODEL_TASK_LEVEL = {
"文本生成": "text-generation",
"视觉多模态理解": "visual-multi-modal",
"图片生成图片": "text-to-image-generation",
"文生图": "text-to-image-generation",
"特征抽取": "feature_emb",
"句子相似度": "feature_emb",
"问答": "question_answering",
"语音识别": "asr",
"视觉分类": "vision_classification",
"零样本图像分类": "vision_classification",
"强化学习": "reinforcement_learning",
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Continuously refill submissions, but only on GPUs with historical verified successes."
@@ -137,21 +107,6 @@ def pipeline_tag_for_task_type(task_type: str | None) -> str | None:
return spec.pipeline_tags[0]
def task_type_from_history_task(task: dict[str, Any]) -> str | None:
task_type = task.get("taskType")
if task_type and str(task_type) in TASK_SPEC_BY_TYPE:
return str(task_type)
level_id = task.get("modelTaskLevelId")
if level_id is not None:
mapped = TASK_TYPE_BY_MODEL_TASK_LEVEL_ID.get(str(level_id))
if mapped:
return mapped
level_name = task.get("modelTaskLevel")
if level_name:
return TASK_TYPE_BY_MODEL_TASK_LEVEL.get(str(level_name))
return None
def successful_models_from_tasks(tasks: list[dict[str, Any]], *, limit: int) -> tuple[list[HFModelSummary], dict[str, set[str]]]:
by_model: dict[str, HFModelSummary] = {}
gpus_by_model: dict[str, set[str]] = {}