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