339 lines
15 KiB
Python
339 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
from common import parse_datetime, utc_now, write_json, write_jsonl
|
|
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log
|
|
from history_stats import is_success, update_history_archive
|
|
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 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."
|
|
)
|
|
parser.add_argument("--history-days", type=int, default=3650, help="UTC lookback window for pulling task history")
|
|
parser.add_argument("--history-begin", help="Override history begin time, for example 2026-01-01 or 2026-01-01T00:00:00Z")
|
|
parser.add_argument("--history-end", help="Override history end time; defaults to now")
|
|
parser.add_argument("--history-page-size", type=int, default=100, help="ModelHub task history page size")
|
|
parser.add_argument("--daily-target", type=int, default=0, help="Total submissions to aim for per UTC day; 0 means unlimited")
|
|
parser.add_argument("--gpu", help="Single GPU alias or platform name, for example: k100")
|
|
parser.add_argument("--gpus", help="Comma-separated GPU aliases/platform names")
|
|
parser.add_argument(
|
|
"--successful-gpus-only",
|
|
action="store_true",
|
|
help="Accepted for backward compatibility; this runner always restricts to historically successful GPUs.",
|
|
)
|
|
parser.add_argument("--min-downloads", type=int, default=50, help="Minimum ModelScope download threshold")
|
|
parser.add_argument(
|
|
"--history-stats-threshold",
|
|
type=int,
|
|
default=500,
|
|
help="Minimum local ledger records before using history stats to rank submissions",
|
|
)
|
|
parser.add_argument("--read-concurrency", type=int, default=4, help="Concurrency for read-only remote calls")
|
|
parser.add_argument("--dry-run", action="store_true", help="Plan the queue without creating tasks")
|
|
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
|
parser.add_argument("--runs-dir", default="runs", help=argparse.SUPPRESS)
|
|
parser.add_argument("--ledger-path", default="ledger/submissions.jsonl", help=argparse.SUPPRESS)
|
|
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
|
parser.add_argument("--history-archive-limit", type=int, default=20000, help=argparse.SUPPRESS)
|
|
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
|
parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS)
|
|
parser.add_argument("--outcomes-path", default="outcomes/submissions.jsonl", help=argparse.SUPPRESS)
|
|
parser.add_argument("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--modelhub-base-url", default=os.getenv("MODELHUB_BASE_URL", "https://modelhub.org.cn"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--modelhub-token", default=os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--modelscope-token", default=os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN"), help=argparse.SUPPRESS)
|
|
parser.add_argument("--poll-interval-seconds", type=int, default=60, help="Sleep between polling cycles when no slots are available")
|
|
parser.add_argument("--idle-interval-seconds", type=int, default=30, help="Sleep between cycles when a scan submits nothing")
|
|
parser.add_argument(
|
|
"--post-cycle-cooldown-seconds",
|
|
type=int,
|
|
default=1,
|
|
help="Short sleep after a successful cycle to avoid same-second run-dir collisions",
|
|
)
|
|
parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached")
|
|
return parser
|
|
|
|
|
|
def make_modelhub_client(args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool:
|
|
tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
|
|
if len(tokens) > 1:
|
|
return ModelHubClientPool([ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in tokens])
|
|
return ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url)
|
|
|
|
|
|
def history_window(args: argparse.Namespace, now) -> tuple[Any, Any]:
|
|
begin = parse_datetime(args.history_begin) if args.history_begin else now - timedelta(days=args.history_days)
|
|
end = parse_datetime(args.history_end) if args.history_end else now
|
|
if begin is None:
|
|
raise ValueError(f"Invalid --history-begin value: {args.history_begin}")
|
|
if end is None:
|
|
raise ValueError(f"Invalid --history-end value: {args.history_end}")
|
|
return begin, end
|
|
|
|
|
|
def model_id_from_task(task: dict[str, Any]) -> str | None:
|
|
value = task.get("modelId") or task.get("modelName") or task.get("model")
|
|
if value:
|
|
return str(value).strip()
|
|
address = str(task.get("modelAddress") or "").strip().rstrip("/")
|
|
marker = "modelscope.cn/models/"
|
|
if marker in address:
|
|
return address.split(marker, 1)[1].strip("/")
|
|
marker = "huggingface.co/"
|
|
if marker in address:
|
|
return address.split(marker, 1)[1].strip("/")
|
|
return None
|
|
|
|
|
|
def pipeline_tag_for_task_type(task_type: str | None) -> str | None:
|
|
if not task_type:
|
|
return None
|
|
spec = TASK_SPEC_BY_TYPE.get(task_type)
|
|
if spec is None:
|
|
return 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]] = {}
|
|
for task in tasks:
|
|
if not is_success(task):
|
|
continue
|
|
model_id = model_id_from_task(task)
|
|
if not model_id:
|
|
continue
|
|
task_type = task_type_from_history_task(task)
|
|
pipeline_tag = pipeline_tag_for_task_type(task_type)
|
|
if not pipeline_tag:
|
|
continue
|
|
timestamp = parse_datetime(task.get("updateTime")) or parse_datetime(task.get("createTime")) or parse_datetime(task.get("submitTime"))
|
|
current = by_model.get(model_id)
|
|
if current is None or (timestamp and (current.last_modified is None or timestamp > current.last_modified)):
|
|
by_model[model_id] = HFModelSummary(
|
|
repo_id=model_id,
|
|
downloads=0,
|
|
last_modified=timestamp,
|
|
pipeline_tag=pipeline_tag,
|
|
)
|
|
gpu_type = task.get("gpuType") or task.get("targetGpu")
|
|
if gpu_type:
|
|
gpus_by_model.setdefault(model_id, set()).add(str(gpu_type))
|
|
|
|
models = sorted(by_model.values(), key=lambda item: item.last_modified or parse_datetime("1970-01-01"), reverse=True)
|
|
if limit > 0:
|
|
models = models[:limit]
|
|
return models, {model.repo_id: gpus_by_model.get(model.repo_id, set()) for model in models}
|
|
|
|
|
|
def successful_gpus_from_tasks(tasks: list[dict[str, Any]]) -> list[str]:
|
|
latest_success_by_gpu: dict[str, Any] = {}
|
|
for task in tasks:
|
|
if not is_success(task):
|
|
continue
|
|
gpu_type = task.get("gpuType") or task.get("targetGpu")
|
|
if not gpu_type:
|
|
continue
|
|
gpu = str(gpu_type)
|
|
timestamp = parse_datetime(task.get("updateTime")) or parse_datetime(task.get("createTime")) or parse_datetime(task.get("submitTime"))
|
|
current = latest_success_by_gpu.get(gpu)
|
|
if current is None or (timestamp and (current is None or timestamp > current)):
|
|
latest_success_by_gpu[gpu] = timestamp
|
|
return [
|
|
gpu for gpu, _ in sorted(
|
|
latest_success_by_gpu.items(),
|
|
key=lambda item: ((item[1].timestamp() if item[1] else 0.0), item[0]),
|
|
reverse=True,
|
|
)
|
|
]
|
|
|
|
|
|
def resolve_successful_target_gpus(
|
|
args: argparse.Namespace,
|
|
successful_gpus: list[str],
|
|
template_selector: TemplateSelector,
|
|
) -> list[str]:
|
|
normalized_successful: list[str] = []
|
|
for gpu in successful_gpus:
|
|
normalized = template_selector.normalize_gpu(gpu)
|
|
if normalized not in normalized_successful:
|
|
normalized_successful.append(normalized)
|
|
|
|
requested_values: list[str] = []
|
|
if getattr(args, "gpus", None):
|
|
requested_values.extend(value.strip() for value in str(args.gpus).split(",") if value.strip())
|
|
if getattr(args, "gpu", None):
|
|
requested_values.append(str(args.gpu).strip())
|
|
|
|
if requested_values:
|
|
requested_normalized: list[str] = []
|
|
for value in requested_values:
|
|
normalized = template_selector.normalize_gpu(value)
|
|
if normalized not in requested_normalized:
|
|
requested_normalized.append(normalized)
|
|
filtered = [gpu for gpu in requested_normalized if gpu in normalized_successful]
|
|
else:
|
|
filtered = normalized_successful
|
|
|
|
if not filtered:
|
|
raise RuntimeError("No historically successful GPUs are available after applying the requested GPU filter")
|
|
return filtered
|
|
|
|
|
|
def make_poll_args(args: argparse.Namespace, target_gpus: list[str]) -> argparse.Namespace:
|
|
poll_args = argparse.Namespace(**vars(args))
|
|
poll_args.gpu = None
|
|
poll_args.gpus = ",".join(target_gpus)
|
|
return poll_args
|
|
|
|
|
|
def run_success_history_poll(
|
|
args: argparse.Namespace,
|
|
*,
|
|
now=None,
|
|
run_fn: Callable[..., dict[str, Any]] = run_poll_loop,
|
|
modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
|
|
template_selector: TemplateSelector | None = None,
|
|
) -> dict[str, Any]:
|
|
now = now or utc_now()
|
|
template_selector = template_selector or TemplateSelector()
|
|
modelhub_client = modelhub_client or make_modelhub_client(args)
|
|
|
|
begin, end = history_window(args, now)
|
|
tasks = modelhub_client.list_tasks(page_size=args.history_page_size, only_mine=True, begin_time=begin, end_time=end)
|
|
archived_history = update_history_archive(Path(args.history_archive_path), tasks, limit=args.history_archive_limit)
|
|
|
|
successful_models, gpus_by_model = successful_models_from_tasks(tasks, limit=0)
|
|
successful_gpus = successful_gpus_from_tasks(tasks)
|
|
target_gpus = resolve_successful_target_gpus(args, successful_gpus, template_selector)
|
|
poll_args = make_poll_args(args, target_gpus)
|
|
|
|
summary = run_fn(
|
|
base_args=poll_args,
|
|
now=now,
|
|
modelhub_client=modelhub_client,
|
|
template_selector=template_selector,
|
|
)
|
|
|
|
poll_run_dir = Path(summary["pollRunDir"])
|
|
write_json(
|
|
poll_run_dir / "history_success_gpus.json",
|
|
{
|
|
"generatedAt": now.isoformat(),
|
|
"historyBegin": begin.isoformat(),
|
|
"historyEnd": end.isoformat(),
|
|
"historyRecordCount": len(tasks),
|
|
"successfulGpuCount": len(successful_gpus),
|
|
"successfulGpus": successful_gpus,
|
|
"filteredTargetGpus": target_gpus,
|
|
},
|
|
)
|
|
write_jsonl(
|
|
poll_run_dir / "history_success_models.jsonl",
|
|
[
|
|
{
|
|
"repoId": model.repo_id,
|
|
"pipelineTag": model.pipeline_tag,
|
|
"lastSuccessfulUpdate": model.last_modified.isoformat() if model.last_modified else None,
|
|
"successfulGpus": sorted(gpus_by_model.get(model.repo_id, set())),
|
|
}
|
|
for model in successful_models
|
|
],
|
|
)
|
|
|
|
summary.update(
|
|
{
|
|
"historyBegin": begin.isoformat(),
|
|
"historyEnd": end.isoformat(),
|
|
"historyRecordCount": len(tasks),
|
|
"historyArchiveRecordCount": len(archived_history),
|
|
"successfulGpuCount": len(successful_gpus),
|
|
"successfulGpus": successful_gpus,
|
|
"filteredTargetGpus": target_gpus,
|
|
"successfulModelCount": len(successful_models),
|
|
"successfulGpuAuditPath": str(poll_run_dir / "history_success_gpus.json"),
|
|
"successfulModelAuditPath": str(poll_run_dir / "history_success_models.jsonl"),
|
|
}
|
|
)
|
|
write_json(poll_run_dir / "summary.json", summary)
|
|
return summary
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
ensure_tokens(args)
|
|
log(
|
|
f"[success-history] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} "
|
|
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'} "
|
|
f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
|
|
)
|
|
summary = run_success_history_poll(args)
|
|
print(f"poll_run_dir={summary['pollRunDir']}")
|
|
print(f"history_records={summary['historyRecordCount']} successful_gpus={summary['successfulGpuCount']} successful_models={summary['successfulModelCount']}")
|
|
print(f"filtered_target_gpus={','.join(summary['filteredTargetGpus'])}")
|
|
print(f"submitted_total={summary['submittedTotal']} cycles={summary['cycles']} stopped_reason={summary['stoppedReason']}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|