from __future__ import annotations import argparse import os from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import timedelta from pathlib import Path from typing import Any from common import parse_datetime, utc_now, write_json, write_jsonl from hf_discovery import HuggingFaceDiscovery from history_stats import ( append_ledger_entry, build_empty_pre_submit_report, count_submissions_for_day, load_ledger, update_history_archive, ) from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool from models import CandidateModel, HFModelSummary, ModelInspection from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker from task_registry import TASK_SPEC_BY_TYPE, all_task_types, choose_framework_for_task, choose_text_generation_framework, pipeline_tags_for_task_types, task_specs_for_model from template_selector import TemplateSelector DEFAULT_RUNS_DIR = Path("runs") DEFAULT_LEDGER_PATH = Path("ledger/submissions.jsonl") def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Auto discover and submit public ModelScope models to ModelHub.") 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 from templates.") parser.add_argument("--task-types", help="Comma-separated ModelHub task types. Omit to auto-enable all supported task types.") parser.add_argument("--limit", type=int, default=300, help="Maximum number of ModelScope models to scan per pipeline tag") 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 downloads threshold") parser.add_argument("--daily-target", type=int, default=0, help="Daily submission target across all auto runs; 0 means unlimited") parser.add_argument("--dry-run", action="store_true", help="Only write artifacts without creating tasks") 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") group = parser.add_mutually_exclusive_group() group.add_argument("--since-hours", type=int, help="Only scan models updated within the last N hours") group.add_argument("--updated-after", help="Only scan models updated after the given UTC time") parser.add_argument("--stats-window-days", type=int, default=7, help="History window in days") 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("--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("--outcomes-path", default=str(DEFAULT_OUTCOMES_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("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), 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) return parser def determine_updated_after(args: argparse.Namespace, now) -> Any: if args.since_hours is not None: return now - timedelta(hours=args.since_hours) if args.updated_after: parsed = parse_datetime(args.updated_after) if parsed is None: raise ValueError(f"Invalid --updated-after value: {args.updated_after}") return parsed return None def resolve_task_types(args: argparse.Namespace) -> list[str]: if not args.task_types: return all_task_types() selected = [value.strip() for value in args.task_types.split(",") if value.strip()] unknown = [task_type for task_type in selected if task_type not in TASK_SPEC_BY_TYPE] if unknown: raise ValueError(f"Unknown task types: {', '.join(sorted(unknown))}") return selected def resolve_target_gpus(args: argparse.Namespace, selector: TemplateSelector, task_types: list[str]) -> list[str]: raw_values: list[str] = [] if args.gpus: raw_values.extend(value.strip() for value in args.gpus.split(",") if value.strip()) if args.gpu: raw_values.append(args.gpu) if raw_values: normalized: list[str] = [] for value in raw_values: target_gpu = selector.normalize_gpu(value) if target_gpu not in normalized: normalized.append(target_gpu) return normalized discovered: list[str] = [] for task_type in task_types: for target_gpu in selector.supported_target_gpus(task_type, auto_only=True): if target_gpu not in discovered: discovered.append(target_gpu) return discovered def resolve_scan_limit( args: argparse.Namespace, *, target_gpu_count: int, remaining_daily_quota: int, platform_available_slots: int | None, ) -> int: base_limit = max(1, args.limit) explicit_max = max(0, int(getattr(args, "max_scan_models", 0) or 0)) if explicit_max > 0: return min(base_limit, explicit_max) multiplier = max(1, int(getattr(args, "scan_multiplier", 4) or 1)) if args.daily_target > 0: needed = max(1, remaining_daily_quota) elif platform_available_slots is not None and platform_available_slots > 0: needed = max(1, platform_available_slots) else: needed = max(200, target_gpu_count * 20) max_submits_per_run = max(0, int(getattr(args, "max_submits_per_run", 0) or 0)) if max_submits_per_run > 0: needed = min(needed, max_submits_per_run) return min(base_limit, needed * multiplier) def choose_candidate_for_gpu( *, model: HFModelSummary, inspection: ModelInspection, template_selector: TemplateSelector, task_types: list[str], target_gpu: str, ) -> CandidateModel | None: for task_type in task_types: supported_frameworks = template_selector.supported_frameworks_for_auto(task_type, target_gpu) if not supported_frameworks: continue try: framework = choose_framework_for_task(task_type, target_gpu, supported_frameworks, inspection) template = template_selector.select_template(task_type, framework, target_gpu) config_params = template_selector.render_config( template, gguf_filename=inspection.selected_gguf if framework == "llamacpp" else None, ) except ValueError: continue spec = TASK_SPEC_BY_TYPE[task_type] return CandidateModel( repo_id=model.repo_id, model_address=model.model_address, pipeline_tag=model.pipeline_tag, modality=spec.modality, task_type=task_type, target_gpu=target_gpu, framework=framework, template_id=template.template_id, config_params=config_params, downloads=model.downloads, last_modified=model.last_modified, gguf_filename=inspection.selected_gguf, score=0.0, warnings=[], ) return None def resolve_submit_concurrency( args: argparse.Namespace, *, modelhub_client: ModelHubClient | ModelHubClientPool | Any, planned_submit_count: int, ) -> int: explicit_workers = int(getattr(args, "submit_concurrency", 0) or 0) if explicit_workers > 0: return explicit_workers clients = getattr(modelhub_client, "clients", None) if isinstance(clients, list) and clients: if planned_submit_count > 0: return max(1, min(len(clients), planned_submit_count)) return len(clients) return 1 def resolve_max_submit_count( args: argparse.Namespace, planned_count: int, remaining_daily_quota: int, ) -> int: explicit_limit = int(getattr(args, "max_submits_per_run", 0) or 0) planned_submit_count = planned_count if remaining_daily_quota > 0: planned_submit_count = min(planned_submit_count, remaining_daily_quota) if explicit_limit > 0: planned_submit_count = min(planned_submit_count, explicit_limit) return planned_submit_count def process_model_for_candidates( *, model: HFModelSummary, hf_discovery: HuggingFaceDiscovery, modelhub_client: ModelHubClient, template_selector: TemplateSelector, target_gpus: list[str], allowed_task_types: list[str], outcome_tracker: OutcomeTracker | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]: specs = [spec for spec in task_specs_for_model(model) if spec.task_type in allowed_task_types] if not specs: return [], [{"repoId": model.repo_id, "reason": f"unsupported_pipeline_tag:{model.pipeline_tag or 'unknown'}"}], [] processed_gpus = modelhub_client.processed_gpus_for_model(model.repo_id) candidates: list[dict[str, Any]] = [] skipped: list[dict[str, Any]] = [] failed: list[dict[str, Any]] = [] allowed_task_types_set = [spec.task_type for spec in specs] pending_task_types_by_gpu: list[tuple[str, list[str]]] = [] for target_gpu in target_gpus: if target_gpu in processed_gpus: skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "already_processed_for_gpu"}) continue if outcome_tracker and outcome_tracker.is_model_gpu_failed(model.repo_id, target_gpu): skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "previously_failed_locally"}) continue compatible_task_types = [ task_type for task_type in allowed_task_types_set if template_selector.supported_frameworks_for_auto(task_type, target_gpu) ] if not compatible_task_types: skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "no_compatible_auto_template_or_framework"}) continue pending_task_types_by_gpu.append((target_gpu, compatible_task_types)) if not pending_task_types_by_gpu: return candidates, skipped, failed try: inspection = hf_discovery.inspect_model(model) except Exception as exc: return [], skipped, [{"repoId": model.repo_id, "reason": str(exc)}] for target_gpu, task_types in pending_task_types_by_gpu: best = choose_candidate_for_gpu( model=model, inspection=inspection, template_selector=template_selector, task_types=task_types, target_gpu=target_gpu, ) if best is None: skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "no_compatible_auto_template_or_framework"}) continue candidates.append(candidate_to_record(best)) return candidates, skipped, failed def submit_candidate( candidate: dict[str, Any], modelhub_client: ModelHubClient, ) -> dict[str, Any]: payload = { "modelAddress": candidate["modelAddress"], "taskType": candidate["taskType"], "targetGpu": candidate["targetGpu"], "framework": candidate["framework"], "configParams": candidate["configParams"], } strategy_id = os.getenv("STRATEGY_ID") if strategy_id: payload["strategyId"] = strategy_id submit_time = utc_now() try: response = modelhub_client.add_task(payload) task_id = extract_task_id_from_submit_response(response) if task_id is None: task_id = modelhub_client.find_recent_task_id(candidate["repoId"], candidate["targetGpu"], submit_time) return { "outcome": "submitted", "candidate": candidate, "submitTime": submit_time.isoformat(), "taskId": task_id, "responseData": response.get("data"), } except ModelHubAPIError as exc: print( f"[submit] failed repo={candidate['repoId']} gpu={candidate['targetGpu']} " f"framework={candidate['framework']} reason={exc}", flush=True, ) return { "outcome": "failed", "candidate": candidate, "reason": str(exc), } def make_run_dir(runs_dir: Path, now) -> Path: base_name = now.strftime("%Y%m%dT%H%M%SZ") run_dir = runs_dir / base_name if not run_dir.exists(): return run_dir for suffix in range(1, 1000): candidate = runs_dir / f"{base_name}.{suffix:03d}" if not candidate.exists(): return candidate return candidate def run_submission( args: argparse.Namespace, *, now=None, 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() template_selector = template_selector or TemplateSelector() selected_task_types = resolve_task_types(args) target_gpus = resolve_target_gpus(args, template_selector, selected_task_types) if not target_gpus: raise RuntimeError("No auto-submittable GPUs are available for the selected task types") hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=args.hf_base_url) if modelhub_client is None: modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token])) if len(modelhub_tokens) > 1: clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in modelhub_tokens] modelhub_client = ModelHubClientPool(clients) else: modelhub_client = ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url) runs_dir = Path(args.runs_dir) ledger_path = Path(args.ledger_path) history_archive_path = Path(args.history_archive_path) run_dir = make_run_dir(runs_dir, now) run_dir.mkdir(parents=True, exist_ok=True) ledger_path.parent.mkdir(parents=True, exist_ok=True) history_archive_path.parent.mkdir(parents=True, exist_ok=True) outcome_tracker = outcome_tracker or OutcomeTracker(Path(args.outcomes_path)) synced_count = 0 if not getattr(args, "skip_outcome_sync", False): try: synced_count = outcome_tracker.sync_from_api(modelhub_client) except Exception: pass updated_after = determine_updated_after(args, now) history_begin = now - timedelta(days=args.stats_window_days) ledger_entries = load_ledger(ledger_path) day_start = now.replace(hour=0, minute=0, second=0, microsecond=0) # Count today's submissions from the local ledger (avoids expensive paginated API call) daily_snapshot = count_submissions_for_day(tasks=[], ledger_entries=ledger_entries, day_start=day_start, day_end=now) # Fallback: if ledger has no entries yet, do a quick API check (limit to 1 page) if daily_snapshot["totalCount"] <= 0: today_tasks_page = modelhub_client.list_tasks_page( current=1, page_size=20, only_mine=True, begin_time=day_start, end_time=now, ) today_tasks = (today_tasks_page.get("data") or {}).get("records") or [] daily_snapshot = count_submissions_for_day(tasks=today_tasks, ledger_entries=ledger_entries, day_start=day_start, day_end=now) platform_available_slots = None if hasattr(modelhub_client, "available_submit_slots"): platform_available_slots = max(0, int(modelhub_client.available_submit_slots())) unlimited_daily_target = args.daily_target <= 0 if unlimited_daily_target: remaining_daily_quota = platform_available_slots if platform_available_slots is not None else 1_000_000_000 else: remaining_daily_quota = max(0, args.daily_target - daily_snapshot["totalCount"]) if platform_available_slots is not None: remaining_daily_quota = min(remaining_daily_quota, platform_available_slots) history_report_reason = "history_archive_skipped" if getattr(args, "skip_history_archive", False) else "history_archive_only_mode" if args.daily_target > 0 and remaining_daily_quota <= 0: archived_history: list[dict[str, Any]] = [] report = build_empty_pre_submit_report( window_days=args.stats_window_days, generated_at=now, reason="daily_target_reached_before_scan", ledger_entries=len(ledger_entries), ) write_json(run_dir / "pre_submit_report.json", report) return { "generatedAt": now.isoformat(), "dryRun": bool(args.dry_run), "selectedTaskTypes": selected_task_types, "targetGpus": target_gpus, "dailyTarget": args.daily_target, "unlimitedDailyTarget": unlimited_daily_target, "submittedTodayBeforeRun": daily_snapshot["totalCount"], "remainingDailyQuotaBeforeRun": remaining_daily_quota, "platformAvailableSlotsBeforeRun": platform_available_slots, "historyStatsEnabled": False, "historyStatsThreshold": getattr(args, "history_stats_threshold", 500), "historyArchivePath": str(history_archive_path), "historyArchiveRecordCount": len(archived_history), "scanLimit": 0, "scannedModels": 0, "candidateCount": 0, "plannedSubmitCount": 0, "submittedCount": 0, "skippedCount": 0, "failedCount": 0, "warnings": report.get("warnings", []), "runDir": str(run_dir), "outcomeSyncCount": synced_count, } if getattr(args, "skip_history_archive", False): archived_history = [] else: history_tasks = modelhub_client.list_tasks( page_size=50, only_mine=True, begin_time=history_begin, end_time=now, ) archived_history = update_history_archive( history_archive_path, history_tasks, limit=getattr(args, "history_archive_limit", 5000), ) report = build_empty_pre_submit_report( window_days=args.stats_window_days, generated_at=now, reason=history_report_reason, ledger_entries=len(ledger_entries), ) write_json(run_dir / "pre_submit_report.json", report) scan_limit = resolve_scan_limit( args, target_gpu_count=len(target_gpus), remaining_daily_quota=remaining_daily_quota, platform_available_slots=platform_available_slots, ) model_query_kwargs = { "pipeline_tags": pipeline_tags_for_task_types(selected_task_types), "limit": scan_limit, "min_downloads": args.min_downloads, "updated_after": updated_after, } try: model_query_kwargs["read_concurrency"] = max(1, args.read_concurrency) models = hf_discovery.list_recent_models(**model_query_kwargs) except TypeError: models = hf_discovery.list_recent_models( pipeline_tags=model_query_kwargs["pipeline_tags"], limit=model_query_kwargs["limit"], min_downloads=model_query_kwargs["min_downloads"], updated_after=model_query_kwargs["updated_after"], ) candidates: list[dict[str, Any]] = [] skipped: list[dict[str, Any]] = [] failed: list[dict[str, Any]] = [] with ThreadPoolExecutor(max_workers=max(1, args.read_concurrency)) as executor: futures = { executor.submit( process_model_for_candidates, model=model, hf_discovery=hf_discovery, modelhub_client=modelhub_client, template_selector=template_selector, target_gpus=target_gpus, allowed_task_types=selected_task_types, outcome_tracker=outcome_tracker, ): index for index, model in enumerate(models) } ordered_results: dict[int, tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]] = {} for future in as_completed(futures): index = futures[future] model = models[index] try: ordered_results[index] = future.result() except Exception as exc: ordered_results[index] = ([], [], [{"repoId": model.repo_id, "reason": str(exc)}]) for index in range(len(models)): model_candidates, model_skipped, model_failed = ordered_results.get(index, ([], [], [])) candidates.extend(model_candidates) skipped.extend(model_skipped) failed.extend(model_failed) write_jsonl(run_dir / "candidates.jsonl", candidates) submitted: list[dict[str, Any]] = [] planned_submit_count = resolve_max_submit_count( args=args, planned_count=len(candidates), remaining_daily_quota=remaining_daily_quota, ) planned_candidates = candidates[:planned_submit_count] submit_workers = 1 if not args.dry_run: submit_workers = resolve_submit_concurrency( args, modelhub_client=modelhub_client, planned_submit_count=len(planned_candidates), ) with ThreadPoolExecutor(max_workers=submit_workers) as executor: futures = { executor.submit( submit_candidate, candidate, modelhub_client, ): index for index, candidate in enumerate(planned_candidates) } ordered_results: dict[int, dict[str, Any]] = {} for future in as_completed(futures): index = futures[future] try: ordered_results[index] = future.result() except Exception as exc: candidate = planned_candidates[index] ordered_results[index] = { "outcome": "failed", "candidate": candidate, "reason": str(exc), } for index in range(len(planned_candidates)): result = ordered_results.get(index) if result is None: continue candidate = result["candidate"] if result["outcome"] == "failed": failed.append( { "repoId": candidate["repoId"], "targetGpu": candidate["targetGpu"], "framework": candidate["framework"], "taskType": candidate["taskType"], "reason": result.get("reason", "submission_failed"), } ) continue submitted_record = { **candidate, "submitTime": result["submitTime"], "taskId": result["taskId"], "responseData": result["responseData"], } submitted.append(submitted_record) append_ledger_entry( ledger_path, { "modelId": candidate["repoId"], "modelAddress": candidate["modelAddress"], "targetGpu": candidate["targetGpu"], "framework": candidate["framework"], "templateId": candidate["templateId"], "taskId": result["taskId"], "taskType": candidate["taskType"], "submitTime": result["submitTime"], }, ) outcome_tracker.record_submission( model_id=candidate["repoId"], target_gpu=candidate["targetGpu"], framework=candidate["framework"], task_type=candidate["taskType"], task_id=result["taskId"], submit_time=result["submitTime"], ) write_jsonl(run_dir / "submitted.jsonl", submitted) write_jsonl(run_dir / "skipped.jsonl", skipped) write_jsonl(run_dir / "failed.jsonl", failed) outcome_tracker.save() summary = { "generatedAt": now.isoformat(), "dryRun": bool(args.dry_run), "selectedTaskTypes": selected_task_types, "targetGpus": target_gpus, "maxSubmitsPerRun": int(getattr(args, "max_submits_per_run", 0) or 0), "dailyTarget": args.daily_target, "unlimitedDailyTarget": unlimited_daily_target, "submittedTodayBeforeRun": daily_snapshot["totalCount"], "remainingDailyQuotaBeforeRun": remaining_daily_quota, "platformAvailableSlotsBeforeRun": platform_available_slots, "historyStatsEnabled": False, "historyStatsThreshold": getattr(args, "history_stats_threshold", 500), "historyArchivePath": str(history_archive_path), "historyArchiveRecordCount": len(archived_history), "scanLimit": scan_limit, "scannedModels": len(models), "candidateCount": len(candidates), "plannedSubmitCount": len(planned_candidates), "submittedCount": len(submitted), "skippedCount": len(skipped), "failedCount": len(failed), "submitConcurrencyUsed": submit_workers, "warnings": report.get("warnings", []), "runDir": str(run_dir), "outcomeSyncCount": synced_count, } write_json(run_dir / "summary.json", summary) return summary def candidate_to_record(candidate: CandidateModel) -> dict[str, Any]: return { "repoId": candidate.repo_id, "modelAddress": candidate.model_address, "pipelineTag": candidate.pipeline_tag, "modality": candidate.modality, "taskType": candidate.task_type, "taskPriority": TASK_SPEC_BY_TYPE[candidate.task_type].priority, "targetGpu": candidate.target_gpu, "framework": candidate.framework, "templateId": candidate.template_id, "configParams": candidate.config_params, "downloads": candidate.downloads, "lastModified": candidate.last_modified.isoformat() if candidate.last_modified else None, "ggufFilename": candidate.gguf_filename, "score": candidate.score, "warnings": candidate.warnings, } def extract_task_id_from_submit_response(response: dict[str, Any]) -> str | None: data = response.get("data") if isinstance(data, dict) and data.get("id") is not None: return str(data["id"]) return None def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) summary = run_submission(args) print(f"run_dir={summary['runDir']}") print(f"daily_target={summary['dailyTarget']} remaining_before_run={summary['remainingDailyQuotaBeforeRun']}") print(f"task_types={','.join(summary['selectedTaskTypes'])}") print(f"gpus={','.join(summary['targetGpus'])}") print( f"candidates={summary['candidateCount']} planned={summary['plannedSubmitCount']} " f"submitted={summary['submittedCount']} skipped={summary['skippedCount']} failed={summary['failedCount']}" ) return 0 if __name__ == "__main__": raise SystemExit(main())