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())