300 lines
13 KiB
Python
300 lines
13 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
from dataclasses import asdict, dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Callable
|
||
|
|
|
||
|
|
from common import utc_now, write_json
|
||
|
|
from hf_discovery import HuggingFaceDiscovery
|
||
|
|
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission
|
||
|
|
from modelhub_client import ModelHubClient
|
||
|
|
from outcome_tracker import OutcomeTracker
|
||
|
|
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||
|
|
from template_selector import TemplateSelector
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_DAILY_RUNS_DIR = Path("daily_runs")
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class WaveSpec:
|
||
|
|
name: str
|
||
|
|
task_types: tuple[str, ...]
|
||
|
|
gpus: str | None = None
|
||
|
|
limit: int = 400
|
||
|
|
since_hours: int = 168
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_WAVES: tuple[WaveSpec, ...] = (
|
||
|
|
WaveSpec("text_generation_core", ("text-generation",), gpus=None, limit=100000, since_hours=48),
|
||
|
|
# WaveSpec("visual_multimodal_core", ("visual-multi-modal",), gpus=None, limit=400, since_hours=336),
|
||
|
|
# WaveSpec("image_generation", ("text-to-image-generation",), gpus=None, limit=250, since_hours=720),
|
||
|
|
# WaveSpec("asr", ("asr",), gpus=None, limit=250, since_hours=720),
|
||
|
|
# WaveSpec("embedding_and_qa", ("feature_emb", "question_answering"), gpus=None, limit=200, since_hours=720),
|
||
|
|
# WaveSpec("classification", ("text_classification", "vision_classification"), gpus=None, limit=200, since_hours=720),
|
||
|
|
# WaveSpec("reinforcement_learning", ("reinforcement_learning",), gpus=None, limit=200, since_hours=720),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def log(message: str) -> None:
|
||
|
|
print(message, flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
def build_parser() -> argparse.ArgumentParser:
|
||
|
|
parser = argparse.ArgumentParser(description="One-click daily ModelHub submission runner.")
|
||
|
|
parser.add_argument("--daily-target", type=int, default=0, help="Total submissions to aim for per UTC day; 0 means unlimited")
|
||
|
|
parser.add_argument("--rounds", type=int, default=3, help="How many wave cycles to run before stopping")
|
||
|
|
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("--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("--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 run (0 means unlimited)",
|
||
|
|
)
|
||
|
|
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("--hf-base-url", default=os.getenv("HF_BASE_URL", "https://huggingface.co"), help=argparse.SUPPRESS)
|
||
|
|
parser.add_argument("--modelhub-base-url", default=os.getenv("MODELHUB_BASE_URL", "https://modelhub.org.cn"), help=argparse.SUPPRESS)
|
||
|
|
parser.add_argument("--modelhub-token", default=os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN"), help=argparse.SUPPRESS)
|
||
|
|
parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN"), help=argparse.SUPPRESS)
|
||
|
|
return parser
|
||
|
|
|
||
|
|
|
||
|
|
def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argparse.Namespace:
|
||
|
|
user_gpu = getattr(base_args, "gpu", None)
|
||
|
|
user_gpus = getattr(base_args, "gpus", None)
|
||
|
|
return argparse.Namespace(
|
||
|
|
gpu=user_gpu,
|
||
|
|
gpus=user_gpus if (user_gpus or user_gpu) else wave.gpus,
|
||
|
|
task_types=",".join(wave.task_types),
|
||
|
|
limit=wave.limit,
|
||
|
|
min_downloads=base_args.min_downloads,
|
||
|
|
daily_target=base_args.daily_target,
|
||
|
|
dry_run=base_args.dry_run,
|
||
|
|
since_hours=wave.since_hours,
|
||
|
|
updated_after=None,
|
||
|
|
stats_window_days=7,
|
||
|
|
history_stats_threshold=base_args.history_stats_threshold,
|
||
|
|
read_concurrency=base_args.read_concurrency,
|
||
|
|
max_scan_models=getattr(base_args, "max_scan_models", 0),
|
||
|
|
scan_multiplier=getattr(base_args, "scan_multiplier", 4),
|
||
|
|
skip_outcome_sync=getattr(base_args, "skip_outcome_sync", False),
|
||
|
|
skip_history_archive=getattr(base_args, "skip_history_archive", False),
|
||
|
|
submit_concurrency=getattr(base_args, "submit_concurrency", 1),
|
||
|
|
max_submits_per_run=getattr(base_args, "max_submits_per_run", 0),
|
||
|
|
runs_dir=base_args.runs_dir,
|
||
|
|
ledger_path=base_args.ledger_path,
|
||
|
|
outcomes_path=getattr(base_args, "outcomes_path", "outcomes/submissions.jsonl"),
|
||
|
|
history_archive_path=base_args.history_archive_path,
|
||
|
|
history_archive_limit=base_args.history_archive_limit,
|
||
|
|
hf_base_url=base_args.hf_base_url,
|
||
|
|
modelhub_base_url=base_args.modelhub_base_url,
|
||
|
|
modelhub_token=base_args.modelhub_token,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def run_daily_batches(
|
||
|
|
*,
|
||
|
|
base_args: argparse.Namespace,
|
||
|
|
waves: tuple[WaveSpec, ...] = DEFAULT_WAVES,
|
||
|
|
now=None,
|
||
|
|
run_fn: Callable[..., dict[str, Any]] = run_submission,
|
||
|
|
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)
|
||
|
|
if modelhub_client is None:
|
||
|
|
modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
||
|
|
template_selector = template_selector or TemplateSelector()
|
||
|
|
|
||
|
|
daily_run_dir = Path(base_args.daily_runs_dir) / now.strftime("%Y%m%dT%H%M%SZ")
|
||
|
|
daily_run_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
log(f"[daily] daily_run_dir={daily_run_dir}")
|
||
|
|
log(
|
||
|
|
f"[daily] target={base_args.daily_target} rounds={base_args.rounds} "
|
||
|
|
f"dry_run={str(bool(base_args.dry_run)).lower()} waves={len(waves)}"
|
||
|
|
)
|
||
|
|
if hasattr(modelhub_client, "active_task_counts"):
|
||
|
|
active_counts = modelhub_client.active_task_counts()
|
||
|
|
log(f"[daily] modelhub_accounts={len(active_counts)} active_counts={','.join(str(count) for count in active_counts)}")
|
||
|
|
|
||
|
|
wave_results: list[dict[str, Any]] = []
|
||
|
|
submitted_total = 0
|
||
|
|
attempted_waves = 0
|
||
|
|
stopped_reason = "max_rounds_reached"
|
||
|
|
|
||
|
|
for round_index in range(1, base_args.rounds + 1):
|
||
|
|
round_submitted = 0
|
||
|
|
log(f"[daily] round={round_index} start")
|
||
|
|
for wave in waves:
|
||
|
|
attempted_waves += 1
|
||
|
|
wave_args = make_wave_namespace(base_args, wave)
|
||
|
|
log(
|
||
|
|
f"[daily] round={round_index} wave={wave.name} "
|
||
|
|
f"tasks={','.join(wave.task_types)} gpus={wave.gpus or 'auto'} "
|
||
|
|
f"limit={wave.limit} since_hours={wave.since_hours}"
|
||
|
|
)
|
||
|
|
summary = run_fn(
|
||
|
|
wave_args,
|
||
|
|
now=utc_now(),
|
||
|
|
hf_discovery=hf_discovery,
|
||
|
|
modelhub_client=modelhub_client,
|
||
|
|
template_selector=template_selector,
|
||
|
|
outcome_tracker=outcome_tracker,
|
||
|
|
)
|
||
|
|
wave_result = {
|
||
|
|
"round": round_index,
|
||
|
|
"wave": asdict(wave),
|
||
|
|
"summary": summary,
|
||
|
|
}
|
||
|
|
wave_results.append(wave_result)
|
||
|
|
submitted_total += summary["submittedCount"]
|
||
|
|
round_submitted += summary["submittedCount"]
|
||
|
|
log(
|
||
|
|
f"[daily] wave_done name={wave.name} "
|
||
|
|
f"candidates={summary['candidateCount']} planned={summary['plannedSubmitCount']} "
|
||
|
|
f"submitted={summary['submittedCount']} skipped={summary['skippedCount']} "
|
||
|
|
f"failed={summary['failedCount']} remaining_before_run={summary['remainingDailyQuotaBeforeRun']}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if summary.get("platformAvailableSlotsBeforeRun") == 0:
|
||
|
|
stopped_reason = "platform_async_cap_reached"
|
||
|
|
log(f"[daily] stop={stopped_reason}")
|
||
|
|
return finalize_daily_run(
|
||
|
|
daily_run_dir=daily_run_dir,
|
||
|
|
now=now,
|
||
|
|
base_args=base_args,
|
||
|
|
wave_results=wave_results,
|
||
|
|
submitted_total=submitted_total,
|
||
|
|
attempted_waves=attempted_waves,
|
||
|
|
stopped_reason=stopped_reason,
|
||
|
|
)
|
||
|
|
|
||
|
|
if base_args.daily_target > 0 and summary["remainingDailyQuotaBeforeRun"] <= 0:
|
||
|
|
stopped_reason = "daily_target_already_reached"
|
||
|
|
log(f"[daily] stop={stopped_reason}")
|
||
|
|
return finalize_daily_run(
|
||
|
|
daily_run_dir=daily_run_dir,
|
||
|
|
now=now,
|
||
|
|
base_args=base_args,
|
||
|
|
wave_results=wave_results,
|
||
|
|
submitted_total=submitted_total,
|
||
|
|
attempted_waves=attempted_waves,
|
||
|
|
stopped_reason=stopped_reason,
|
||
|
|
)
|
||
|
|
if base_args.daily_target > 0 and not base_args.dry_run and summary["plannedSubmitCount"] <= 0:
|
||
|
|
stopped_reason = "daily_target_reached"
|
||
|
|
log(f"[daily] stop={stopped_reason}")
|
||
|
|
return finalize_daily_run(
|
||
|
|
daily_run_dir=daily_run_dir,
|
||
|
|
now=now,
|
||
|
|
base_args=base_args,
|
||
|
|
wave_results=wave_results,
|
||
|
|
submitted_total=submitted_total,
|
||
|
|
attempted_waves=attempted_waves,
|
||
|
|
stopped_reason=stopped_reason,
|
||
|
|
)
|
||
|
|
|
||
|
|
if round_submitted <= 0:
|
||
|
|
stopped_reason = "no_new_submissions_in_round"
|
||
|
|
log(f"[daily] round={round_index} submitted=0 stop={stopped_reason}")
|
||
|
|
break
|
||
|
|
log(f"[daily] round={round_index} submitted={round_submitted}")
|
||
|
|
|
||
|
|
return finalize_daily_run(
|
||
|
|
daily_run_dir=daily_run_dir,
|
||
|
|
now=now,
|
||
|
|
base_args=base_args,
|
||
|
|
wave_results=wave_results,
|
||
|
|
submitted_total=submitted_total,
|
||
|
|
attempted_waves=attempted_waves,
|
||
|
|
stopped_reason=stopped_reason,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def finalize_daily_run(
|
||
|
|
*,
|
||
|
|
daily_run_dir: Path,
|
||
|
|
now,
|
||
|
|
base_args: argparse.Namespace,
|
||
|
|
wave_results: list[dict[str, Any]],
|
||
|
|
submitted_total: int,
|
||
|
|
attempted_waves: int,
|
||
|
|
stopped_reason: str,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
last_wave_summary = wave_results[-1]["summary"] if wave_results else {}
|
||
|
|
summary = {
|
||
|
|
"generatedAt": now.isoformat(),
|
||
|
|
"dryRun": bool(base_args.dry_run),
|
||
|
|
"dailyTarget": base_args.daily_target,
|
||
|
|
"unlimitedDailyTarget": base_args.daily_target <= 0,
|
||
|
|
"rounds": base_args.rounds,
|
||
|
|
"attemptedWaves": attempted_waves,
|
||
|
|
"submittedTotal": submitted_total,
|
||
|
|
"stoppedReason": stopped_reason,
|
||
|
|
"dailyRunDir": str(daily_run_dir),
|
||
|
|
"remainingDailyQuotaBeforeRun": last_wave_summary.get("remainingDailyQuotaBeforeRun"),
|
||
|
|
"platformAvailableSlotsBeforeRun": last_wave_summary.get("platformAvailableSlotsBeforeRun"),
|
||
|
|
"waveResults": wave_results,
|
||
|
|
}
|
||
|
|
write_json(daily_run_dir / "summary.json", summary)
|
||
|
|
log(
|
||
|
|
f"[daily] finished submitted_total={submitted_total} attempted_waves={attempted_waves} "
|
||
|
|
f"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)
|
||
|
|
log(
|
||
|
|
f"[daily] hf_token={'set' if bool(args.hf_token) else 'missing'} "
|
||
|
|
f"xc_token={'set' if bool(args.modelhub_token) else 'missing'}"
|
||
|
|
)
|
||
|
|
summary = run_daily_batches(base_args=args)
|
||
|
|
print(f"daily_run_dir={summary['dailyRunDir']}")
|
||
|
|
print(f"submitted_total={summary['submittedTotal']}")
|
||
|
|
print(f"attempted_waves={summary['attemptedWaves']}")
|
||
|
|
print(f"stopped_reason={summary['stoppedReason']}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|