adapt for ModelHub agent platform
This commit is contained in:
CoolBoy
681
modelhub_submmit_api/main.py
Normal file
681
modelhub_submmit_api/main.py
Normal file
@@ -0,0 +1,681 @@
|
||||
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
|
||||
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 Hugging Face 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 Hugging Face 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("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)
|
||||
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 | 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:
|
||||
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 | 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)
|
||||
modelhub_client = modelhub_client or 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())
|
||||
Reference in New Issue
Block a user