adapt for ModelHub agent platform
This commit is contained in:
CoolBoy
259
modelhub_submmit_api/poll_runner.py
Normal file
259
modelhub_submmit_api/poll_runner.py
Normal file
@@ -0,0 +1,259 @@
|
||||
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
|
||||
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 Hugging Face 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 HF_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://huggingface.co", 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("--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:
|
||||
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 | 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] hf_token={'set' if bool(args.hf_token) else 'missing'} "
|
||||
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'}"
|
||||
)
|
||||
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())
|
||||
Reference in New Issue
Block a user