from __future__ import annotations import argparse import os from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Callable from common import utc_now, write_json from hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission from modelhub_client import ModelHubClient, ModelHubClientPool from outcome_tracker import OutcomeTracker from runner_common import DEFAULT_KEY_PATH, ensure_tokens from template_selector import TemplateSelector DEFAULT_DAILY_RUNS_DIR = Path("daily_runs") @dataclass(frozen=True) class WaveSpec: name: str task_types: tuple[str, ...] gpus: str | None = None limit: int = 400 since_hours: int = 168 DEFAULT_WAVES: tuple[WaveSpec, ...] = ( WaveSpec("text_generation_core", ("text-generation",), gpus=None, limit=100000, since_hours=48), # WaveSpec("visual_multimodal_core", ("visual-multi-modal",), gpus=None, limit=400, since_hours=336), # WaveSpec("image_generation", ("text-to-image-generation",), gpus=None, limit=250, since_hours=720), # WaveSpec("asr", ("asr",), gpus=None, limit=250, since_hours=720), # WaveSpec("embedding_and_qa", ("feature_emb", "question_answering"), gpus=None, limit=200, since_hours=720), # WaveSpec("classification", ("text_classification", "vision_classification"), gpus=None, limit=200, since_hours=720), # WaveSpec("reinforcement_learning", ("reinforcement_learning",), gpus=None, limit=200, since_hours=720), ) def log(message: str) -> None: print(message, flush=True) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="One-click daily ModelHub submission runner.") parser.add_argument("--daily-target", type=int, default=0, help="Total submissions to aim for per UTC day; 0 means unlimited") parser.add_argument("--rounds", type=int, default=3, help="How many wave cycles to run before stopping") 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. Omit to auto-use all safe GPUs.") parser.add_argument("--max-scan-models", type=int, default=0, help="Hard cap on total scanned models (0 means auto)") parser.add_argument( "--scan-multiplier", type=int, default=4, help="Multiplier used when auto-deriving scan limit from quota/queue capacity", ) 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( "--submit-concurrency", type=int, default=0, help="Concurrency for task submission calls (0 = auto based on token/client count)", ) parser.add_argument( "--max-submits-per-run", type=int, default=0, help="Maximum tasks to submit in one run (0 means unlimited)", ) parser.add_argument("--skip-outcome-sync", action="store_true", help="Skip outcome sync from ModelHub before scanning") parser.add_argument("--skip-history-archive", action="store_true", help="Skip historical task archive download for this run") parser.add_argument("--dry-run", action="store_true", help="Plan the day 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=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), 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=5000, help=argparse.SUPPRESS) parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), 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) return parser def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argparse.Namespace: user_gpu = getattr(base_args, "gpu", None) user_gpus = getattr(base_args, "gpus", None) return argparse.Namespace( gpu=user_gpu, gpus=user_gpus if (user_gpus or user_gpu) else wave.gpus, task_types=",".join(wave.task_types), limit=wave.limit, min_downloads=base_args.min_downloads, daily_target=base_args.daily_target, dry_run=base_args.dry_run, since_hours=wave.since_hours, updated_after=None, stats_window_days=7, history_stats_threshold=base_args.history_stats_threshold, read_concurrency=base_args.read_concurrency, max_scan_models=getattr(base_args, "max_scan_models", 0), scan_multiplier=getattr(base_args, "scan_multiplier", 4), skip_outcome_sync=getattr(base_args, "skip_outcome_sync", False), skip_history_archive=getattr(base_args, "skip_history_archive", False), submit_concurrency=getattr(base_args, "submit_concurrency", 1), max_submits_per_run=getattr(base_args, "max_submits_per_run", 0), runs_dir=base_args.runs_dir, ledger_path=base_args.ledger_path, outcomes_path=getattr(base_args, "outcomes_path", "outcomes/submissions.jsonl"), history_archive_path=base_args.history_archive_path, history_archive_limit=base_args.history_archive_limit, hf_base_url=base_args.hf_base_url, modelhub_base_url=base_args.modelhub_base_url, modelhub_token=base_args.modelhub_token, ) def run_daily_batches( *, base_args: argparse.Namespace, waves: tuple[WaveSpec, ...] = DEFAULT_WAVES, now=None, run_fn: Callable[..., dict[str, Any]] = run_submission, hf_discovery: HuggingFaceDiscovery | None = None, modelhub_client: ModelHubClient | None = None, template_selector: TemplateSelector | None = None, outcome_tracker: OutcomeTracker | None = None, ) -> dict[str, Any]: now = now or utc_now() hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url) if modelhub_client is None: modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token])) if len(modelhub_tokens) > 1: clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens] modelhub_client = ModelHubClientPool(clients) else: modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url) template_selector = template_selector or TemplateSelector() daily_run_dir = Path(base_args.daily_runs_dir) / now.strftime("%Y%m%dT%H%M%SZ") daily_run_dir.mkdir(parents=True, exist_ok=True) log(f"[daily] daily_run_dir={daily_run_dir}") log( f"[daily] target={base_args.daily_target} rounds={base_args.rounds} " f"dry_run={str(bool(base_args.dry_run)).lower()} waves={len(waves)}" ) if hasattr(modelhub_client, "active_task_counts"): active_counts = modelhub_client.active_task_counts() log(f"[daily] modelhub_accounts={len(active_counts)} active_counts={','.join(str(count) for count in active_counts)}") wave_results: list[dict[str, Any]] = [] submitted_total = 0 attempted_waves = 0 stopped_reason = "max_rounds_reached" for round_index in range(1, base_args.rounds + 1): round_submitted = 0 log(f"[daily] round={round_index} start") for wave in waves: attempted_waves += 1 wave_args = make_wave_namespace(base_args, wave) log( f"[daily] round={round_index} wave={wave.name} " f"tasks={','.join(wave.task_types)} gpus={wave.gpus or 'auto'} " f"limit={wave.limit} since_hours={wave.since_hours}" ) summary = run_fn( wave_args, now=utc_now(), hf_discovery=hf_discovery, modelhub_client=modelhub_client, template_selector=template_selector, outcome_tracker=outcome_tracker, ) wave_result = { "round": round_index, "wave": asdict(wave), "summary": summary, } wave_results.append(wave_result) submitted_total += summary["submittedCount"] round_submitted += summary["submittedCount"] log( f"[daily] wave_done name={wave.name} " f"candidates={summary['candidateCount']} planned={summary['plannedSubmitCount']} " f"submitted={summary['submittedCount']} skipped={summary['skippedCount']} " f"failed={summary['failedCount']} remaining_before_run={summary['remainingDailyQuotaBeforeRun']}" ) if summary.get("platformAvailableSlotsBeforeRun") == 0: stopped_reason = "platform_async_cap_reached" log(f"[daily] stop={stopped_reason}") return finalize_daily_run( daily_run_dir=daily_run_dir, now=now, base_args=base_args, wave_results=wave_results, submitted_total=submitted_total, attempted_waves=attempted_waves, stopped_reason=stopped_reason, ) if base_args.daily_target > 0 and summary["remainingDailyQuotaBeforeRun"] <= 0: stopped_reason = "daily_target_already_reached" log(f"[daily] stop={stopped_reason}") return finalize_daily_run( daily_run_dir=daily_run_dir, now=now, base_args=base_args, wave_results=wave_results, submitted_total=submitted_total, attempted_waves=attempted_waves, stopped_reason=stopped_reason, ) if base_args.daily_target > 0 and not base_args.dry_run and summary["plannedSubmitCount"] <= 0: stopped_reason = "daily_target_reached" log(f"[daily] stop={stopped_reason}") return finalize_daily_run( daily_run_dir=daily_run_dir, now=now, base_args=base_args, wave_results=wave_results, submitted_total=submitted_total, attempted_waves=attempted_waves, stopped_reason=stopped_reason, ) if round_submitted <= 0: stopped_reason = "no_new_submissions_in_round" log(f"[daily] round={round_index} submitted=0 stop={stopped_reason}") break log(f"[daily] round={round_index} submitted={round_submitted}") return finalize_daily_run( daily_run_dir=daily_run_dir, now=now, base_args=base_args, wave_results=wave_results, submitted_total=submitted_total, attempted_waves=attempted_waves, stopped_reason=stopped_reason, ) def finalize_daily_run( *, daily_run_dir: Path, now, base_args: argparse.Namespace, wave_results: list[dict[str, Any]], submitted_total: int, attempted_waves: int, stopped_reason: str, ) -> dict[str, Any]: last_wave_summary = wave_results[-1]["summary"] if wave_results else {} summary = { "generatedAt": now.isoformat(), "dryRun": bool(base_args.dry_run), "dailyTarget": base_args.daily_target, "unlimitedDailyTarget": base_args.daily_target <= 0, "rounds": base_args.rounds, "attemptedWaves": attempted_waves, "submittedTotal": submitted_total, "stoppedReason": stopped_reason, "dailyRunDir": str(daily_run_dir), "remainingDailyQuotaBeforeRun": last_wave_summary.get("remainingDailyQuotaBeforeRun"), "platformAvailableSlotsBeforeRun": last_wave_summary.get("platformAvailableSlotsBeforeRun"), "waveResults": wave_results, } write_json(daily_run_dir / "summary.json", summary) log( f"[daily] finished submitted_total={submitted_total} attempted_waves={attempted_waves} " f"stopped_reason={stopped_reason}" ) return summary def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) ensure_tokens(args) log( f"[daily] 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_daily_batches(base_args=args) print(f"daily_run_dir={summary['dailyRunDir']}") print(f"submitted_total={summary['submittedTotal']}") print(f"attempted_waves={summary['attemptedWaves']}") print(f"stopped_reason={summary['stoppedReason']}") return 0 if __name__ == "__main__": raise SystemExit(main())