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 hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR from modelhub_client import ModelHubClient, ModelHubClientPool 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("--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 | ModelHubClientPool: 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] return ModelHubClientPool(clients) 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: 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 = 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_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())