submmit/modelhub_submmit_api/daily_runner.py

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