submmit/modelhub_submmit_api/poll_runner.py

from __future__ import annotations

import argparse
import json
import os
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 gpu_strategy import DEFAULT_GPU_STRATEGY_PATH
from hf_discovery import HuggingFaceDiscovery
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir
from market_intelligence import (
    DEFAULT_FETCH_WORKERS,
    DEFAULT_FRAMEWORK_MIN_SAMPLES,
    DEFAULT_FRAMEWORK_REFRESH_SECONDS,
    DEFAULT_MARKET_INTELLIGENCE_PATH,
    DEFAULT_QUEUE_REFRESH_SECONDS,
    DEFAULT_THROUGHPUT_WINDOW_HOURS,
)
from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
from submission_claims import DEFAULT_CLAIMS_PATH
from template_selector import TemplateSelector
from version import AGENT_VERSION


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("--disable-candidate-preflight", action="store_true", help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-endpoint", default=os.getenv("MODELHUB_LLM_CLASSIFIER_ENDPOINT"), help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-model", default=os.getenv("MODELHUB_LLM_CLASSIFIER_MODEL"), help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-api-key", default=os.getenv("MODELHUB_LLM_CLASSIFIER_API_KEY"), help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-timeout-seconds", type=int, default=int(os.getenv("MODELHUB_LLM_CLASSIFIER_TIMEOUT_SECONDS", "20")), help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-min-deny-confidence", type=float, default=float(os.getenv("MODELHUB_LLM_CLASSIFIER_MIN_DENY_CONFIDENCE", "0.85")), help=argparse.SUPPRESS)
    parser.add_argument("--llm-classifier-cache-path", default=os.getenv("MODELHUB_LLM_CLASSIFIER_CACHE_PATH", ".modelhub_state/llm_classifications.json"), help=argparse.SUPPRESS)
    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("--disable-gpu-strategy", action="store_true", help="Disable adaptive 70/30 proven-GPU scheduling")
    parser.add_argument(
        "--disable-market-intelligence",
        action="store_true",
        help="Disable live queue/throughput and public framework statistics",
    )
    parser.add_argument(
        "--gpu-strategy-refresh-submissions",
        type=int,
        default=200,
        help="Recalculate adaptive GPU choices after this many accepted submissions",
    )
    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(
        "--claims-path",
        default=os.getenv("MODELHUB_AGENT_CLAIMS_PATH", str(DEFAULT_CLAIMS_PATH)),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--submission-exclusions-path",
        default=os.getenv("MODELHUB_SUBMISSION_EXCLUSIONS_PATH", ".modelhub_state/submission_exclusions.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=5000, help=argparse.SUPPRESS)
    parser.add_argument(
        "--gpu-strategy-state-path",
        default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)),
        help=argparse.SUPPRESS,
    )
    parser.add_argument("--gpu-strategy-recent-window", type=int, default=1000, help=argparse.SUPPRESS)
    parser.add_argument("--gpu-strategy-min-long-samples", type=int, default=100, help=argparse.SUPPRESS)
    parser.add_argument(
        "--market-intelligence-state-path",
        default=os.getenv("MODELHUB_MARKET_INTELLIGENCE_PATH", str(DEFAULT_MARKET_INTELLIGENCE_PATH)),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--market-queue-refresh-seconds",
        type=int,
        default=int(os.getenv("MODELHUB_MARKET_QUEUE_REFRESH_SECONDS", str(DEFAULT_QUEUE_REFRESH_SECONDS))),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--market-framework-refresh-seconds",
        type=int,
        default=int(os.getenv("MODELHUB_MARKET_FRAMEWORK_REFRESH_SECONDS", str(DEFAULT_FRAMEWORK_REFRESH_SECONDS))),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--market-throughput-window-hours",
        type=int,
        default=int(os.getenv("MODELHUB_MARKET_THROUGHPUT_WINDOW_HOURS", str(DEFAULT_THROUGHPUT_WINDOW_HOURS))),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--market-fetch-workers",
        type=int,
        default=int(os.getenv("MODELHUB_MARKET_FETCH_WORKERS", str(DEFAULT_FETCH_WORKERS))),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--market-framework-min-samples",
        type=int,
        default=int(os.getenv("MODELHUB_MARKET_FRAMEWORK_MIN_SAMPLES", str(DEFAULT_FRAMEWORK_MIN_SAMPLES))),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--capacity-state-path",
        default=os.getenv("MODELHUB_CAPACITY_STATE_PATH", str(DEFAULT_CAPACITY_STATE_PATH)),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--capacity-probe-interval-cycles",
        type=int,
        default=int(os.getenv("MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES", "3")),
        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("--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=15, help="Sleep between polling cycles when no slots are available")
    parser.add_argument("--idle-interval-seconds", type=int, default=60, help="Sleep between cycles when a scan submits nothing")
    parser.add_argument("--post-cycle-cooldown-seconds", type=int, default=2, 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) -> ModelHubClientPool:
    modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
    token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token]
    clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values]
    return ModelHubClientPool(
        clients,
        capacity_probe_interval_cycles=max(0, int(getattr(base_args, "capacity_probe_interval_cycles", 3) or 0)),
        capacity_state_path=Path(getattr(base_args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)),
    )


def run_poll_loop(
    *,
    base_args: argparse.Namespace,
    now=None,
    run_fn: Callable[..., dict[str, Any]] = run_daily_batches,
    hf_discovery: HuggingFaceDiscovery | None = None,
    modelhub_client: ModelHubClient | ModelHubClientPool | 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)
    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 = make_run_dir(poll_runs_dir, now)

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

    log(f"[poll] version={AGENT_VERSION} 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
        if hasattr(modelhub_client, "configure_capacity_probe"):
            modelhub_client.configure_capacity_probe(cycles)
        active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else []
        capacity_limits = modelhub_client.account_capacity_limits() if hasattr(modelhub_client, "account_capacity_limits") else []
        capacity_probe = modelhub_client.capacity_probe_enabled() if hasattr(modelhub_client, "capacity_probe_enabled") else False
        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"capacity_limits={','.join(str(limit) for limit in capacity_limits) if capacity_limits else 'n/a'} "
            f"capacity_probe={'on' if capacity_probe else 'off'} "
            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_args = _make_cycle_args(base_args)
        cycle_args.capacity_probe_already_configured = True
        cycle_summary = run_fn(
            base_args=cycle_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"duplicates={cycle_summary.get('duplicateTotal', 0)} "
            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:
            duplicate_total = int(cycle_summary.get("duplicateTotal", 0) or 0)
            if duplicate_total > 0 and (available_slots is None or available_slots > 0):
                retry_delay = max(1, int(getattr(base_args, "post_cycle_cooldown_seconds", 2) or 2))
                log(f"[poll] cycle={cycles} sleep={retry_delay}s reason=duplicates_need_replacement_candidates")
                time.sleep(retry_delay)
            else:
                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_token={'set' if bool(args.modelhub_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())