submmit/modelhub_submmit_api/poll_runner.py

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any, Callable

from common import utc_now, write_json
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR
from modelhub_client import ModelHubClient, MultiModelHubClient
from modelscope_discovery import ModelScopeDiscovery
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
from template_selector import TemplateSelector


DEFAULT_POLL_RUNS_DIR = Path("poll_runs")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Continuously poll ModelHub queue slots and refill submissions.")
    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. Omit to auto-use all safe 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(
        "--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 cycle (0 means unlimited)",
    )
    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("--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("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS)
    parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
    parser.add_argument("--hf-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS)
    parser.add_argument("--modelscope-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS)
    parser.add_argument("--modelhub-base-url", default="https://modelhub.org.cn", help=argparse.SUPPRESS)
    parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS)
    parser.add_argument("--hf-token", default=None, help=argparse.SUPPRESS)
    parser.add_argument("--modelscope-token", default=None, 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=0, help="Short sleep after a successful cycle")
    parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached")
    parser.add_argument("--print-stats", action="store_true", help="Load outcomes, sync, print stats report, and exit")
    return parser


def _make_cycle_args(base_args: argparse.Namespace) -> argparse.Namespace:
    cycle_args = argparse.Namespace(**vars(base_args))
    cycle_args.rounds = 1
    # Poll runner handles outcome sync on its own schedule (every OUTCOME_SYNC_INTERVAL cycles).
    # Skip the per-cycle sync inside run_submission to avoid redundant paginated API calls.
    cycle_args.skip_outcome_sync = True
    return cycle_args


def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient:
    modelhub_tokens = list(getattr(base_args, "modelhub_tokens", []) or [])
    if len(modelhub_tokens) > 1:
        return MultiModelHubClient(tokens=modelhub_tokens, base_url=base_args.modelhub_base_url)
    return ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)


def run_poll_loop(
    *,
    base_args: argparse.Namespace,
    now=None,
    run_fn: Callable[..., dict[str, Any]] = run_daily_batches,
    hf_discovery: ModelScopeDiscovery | 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()
    discovery_base_url = getattr(base_args, "modelscope_base_url", None) or base_args.hf_base_url
    hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
    modelhub_client = modelhub_client or _build_modelhub_client(base_args)
    template_selector = template_selector or TemplateSelector()

    poll_runs_dir = Path(base_args.poll_runs_dir)
    poll_runs_dir.mkdir(parents=True, exist_ok=True)
    poll_run_dir = poll_runs_dir / now.strftime("%Y%m%dT%H%M%SZ")
    poll_run_dir.mkdir(parents=True, exist_ok=True)

    outcome_tracker = outcome_tracker or OutcomeTracker(Path(base_args.outcomes_path))
    OUTCOME_SYNC_INTERVAL = 3
    STATS_PRINT_INTERVAL = 10

    log(f"[poll] poll_run_dir={poll_run_dir}")
    log(
        f"[poll] target={base_args.daily_target} dry_run={str(bool(base_args.dry_run)).lower()} "
        f"poll_interval={base_args.poll_interval_seconds}s idle_interval={base_args.idle_interval_seconds}s"
    )

    cycle_summaries: list[dict[str, Any]] = []
    submitted_total = 0
    cycles = 0
    stopped_reason = "max_cycles_reached"

    while True:
        if base_args.max_cycles and cycles >= base_args.max_cycles:
            stopped_reason = "max_cycles_reached"
            break

        cycles += 1
        active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else []
        available_slots = modelhub_client.available_submit_slots() if hasattr(modelhub_client, "available_submit_slots") else None
        log(
            f"[poll] cycle={cycles} active_counts={','.join(str(count) for count in active_counts) if active_counts else 'n/a'} "
            f"available_slots={available_slots if available_slots is not None else 'n/a'}"
        )

        if available_slots is not None and available_slots <= 0:
            log(f"[poll] cycle={cycles} sleep={base_args.poll_interval_seconds}s reason=no_available_slots")
            time.sleep(base_args.poll_interval_seconds)
            continue

        cycle_summary = run_fn(
            base_args=_make_cycle_args(base_args),
            now=utc_now(),
            hf_discovery=hf_discovery,
            modelhub_client=modelhub_client,
            template_selector=template_selector,
            outcome_tracker=outcome_tracker,
        )
        cycle_summaries.append(cycle_summary)
        submitted_total += cycle_summary["submittedTotal"]
        remaining_before_run = cycle_summary.get("remainingDailyQuotaBeforeRun")

        log(
            f"[poll] cycle_done submitted_total={cycle_summary['submittedTotal']} "
            f"remaining_before_run={remaining_before_run if remaining_before_run is not None else 'n/a'} "
            f"stop={cycle_summary['stoppedReason']}"
        )

        if base_args.daily_target > 0 and remaining_before_run is not None and remaining_before_run <= 0:
            stopped_reason = "daily_target_already_reached"
            break

        if cycle_summary["submittedTotal"] <= 0:
            log(f"[poll] cycle={cycles} sleep={base_args.idle_interval_seconds}s reason=no_new_submissions")
            time.sleep(base_args.idle_interval_seconds)
            continue

        if cycles % OUTCOME_SYNC_INTERVAL == 0:
            try:
                synced = outcome_tracker.sync_from_api(modelhub_client)
                if synced > 0:
                    log(f"[poll] cycle={cycles} outcome_sync_updated={synced}")
            except Exception as exc:
                log(f"[poll] cycle={cycles} outcome_sync_error={exc}")

        if cycles % STATS_PRINT_INTERVAL == 0:
            try:
                stats = outcome_tracker.get_stats_report()
                totals = stats.get("totals", {})
                log(
                    f"[poll] cycle={cycles} outcome_stats "
                    f"terminal={totals.get('total', 0)} "
                    f"success={totals.get('successCount', 0)} "
                    f"failed={totals.get('failureCount', 0)} "
                    f"success_rate={totals.get('successRate', 0):.3f} "
                    f"failure_rate={totals.get('failureRate', 0):.3f}"
                )
            except Exception as exc:
                log(f"[poll] cycle={cycles} outcome_stats_error={exc}")

        if hasattr(modelhub_client, "available_submit_slots") and modelhub_client.available_submit_slots() <= 0:
            log(f"[poll] cycle={cycles} sleep={base_args.poll_interval_seconds}s reason=queue_refilled")
            time.sleep(base_args.poll_interval_seconds)
        else:
            cooldown = getattr(base_args, "post_cycle_cooldown_seconds", 1)
            if cooldown > 0:
                time.sleep(cooldown)

    try:
        outcome_tracker.sync_from_api(modelhub_client)
    except Exception:
        pass
    outcome_tracker.save()
    try:
        stats_report = outcome_tracker.get_stats_report()
    except Exception:
        stats_report = {}

    summary = {
        "generatedAt": now.isoformat(),
        "dryRun": bool(base_args.dry_run),
        "dailyTarget": base_args.daily_target,
        "unlimitedDailyTarget": base_args.daily_target <= 0,
        "cycles": cycles,
        "submittedTotal": submitted_total,
        "stoppedReason": stopped_reason,
        "pollRunDir": str(poll_run_dir),
        "cycleSummaries": cycle_summaries,
        "outcomeStats": stats_report,
    }
    write_json(poll_run_dir / "summary.json", summary)
    log(f"[poll] finished submitted_total={submitted_total} cycles={cycles} 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)

    if getattr(args, "print_stats", False):
        try:
            modelhub_client = _build_modelhub_client(args)
            tracker = OutcomeTracker(Path(args.outcomes_path))
            tracker.sync_from_api(modelhub_client)
            report = tracker.get_stats_report()
            print(json.dumps(report, ensure_ascii=False, indent=2))
        except Exception as exc:
            print(f"Failed to generate stats report: {exc}", file=sys.stderr)
            return 1
        return 0

    log(
        f"[poll] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} "
        f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
    )
    summary = run_poll_loop(base_args=args)
    print(f"poll_run_dir={summary['pollRunDir']}")
    print(f"submitted_total={summary['submittedTotal']}")
    print(f"cycles={summary['cycles']}")
    print(f"stopped_reason={summary['stoppedReason']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())