feat: add adaptive GPU scheduling

This commit is contained in:
CoolBoy
2026-08-02 16:59:44 +08:00
parent 80a1b8518d
commit eab5ab6dce
14 changed files with 1216 additions and 52 deletions

View File

@@ -8,6 +8,7 @@ from pathlib import Path
from typing import Any
from common import parse_datetime, runtime_instance_id, utc_now, write_json, write_jsonl
from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH, GPUStrategyManager
from hf_discovery import HuggingFaceDiscovery
from history_stats import (
append_ledger_entry,
@@ -16,7 +17,7 @@ from history_stats import (
load_ledger,
update_history_archive,
)
from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool, is_duplicate_submission_error
from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubAPIError, ModelHubClient, ModelHubClientPool, is_duplicate_submission_error
from models import CandidateModel, HFModelSummary, ModelInspection
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
from submission_claims import DEFAULT_CLAIMS_PATH, SubmissionClaimStore, candidate_key, diversify_candidates
@@ -44,6 +45,17 @@ def build_parser() -> argparse.ArgumentParser:
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(
"--disable-gpu-strategy",
action="store_true",
help="Disable adaptive 50/30/20 GPU scheduling and keep the legacy candidate order",
)
parser.add_argument(
"--gpu-strategy-refresh-submissions",
type=int,
default=200,
help="Recalculate adaptive GPU choices after this many accepted submissions",
)
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()
@@ -79,6 +91,24 @@ def build_parser() -> argparse.ArgumentParser:
)
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(
"--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("--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)
@@ -376,7 +406,11 @@ def run_submission(
modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
token_values: list[str | None] = modelhub_tokens or [args.modelhub_token]
clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in token_values]
modelhub_client = ModelHubClientPool(clients)
modelhub_client = ModelHubClientPool(
clients,
capacity_probe_interval_cycles=max(0, int(getattr(args, "capacity_probe_interval_cycles", 3) or 0)),
capacity_state_path=Path(getattr(args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)),
)
if hasattr(modelhub_client, "begin_cycle"):
modelhub_client.begin_cycle()
@@ -389,6 +423,13 @@ def run_submission(
history_archive_path.parent.mkdir(parents=True, exist_ok=True)
outcome_tracker = outcome_tracker or OutcomeTracker(Path(args.outcomes_path))
strategy_manager: GPUStrategyManager | None = None
strategy_summary: dict[str, Any] = {"enabled": False}
strategy_enabled = (
not bool(getattr(args, "disable_gpu_strategy", False))
and not bool(getattr(args, "gpu", None) or getattr(args, "gpus", None))
and len(target_gpus) > 1
)
synced_count = 0
if not getattr(args, "skip_outcome_sync", False):
try:
@@ -466,6 +507,7 @@ def run_submission(
"historyStatsThreshold": getattr(args, "history_stats_threshold", 500),
"historyArchivePath": str(history_archive_path),
"historyArchiveRecordCount": len(archived_history),
"gpuStrategy": strategy_summary,
"scanLimit": 0,
"scannedModels": 0,
"candidateCount": 0,
@@ -481,6 +523,16 @@ def run_submission(
"outcomeSyncCount": synced_count,
}
if strategy_enabled:
strategy_manager = GPUStrategyManager(
Path(getattr(args, "gpu_strategy_state_path", DEFAULT_GPU_STRATEGY_PATH)),
refresh_submissions=max(1, int(getattr(args, "gpu_strategy_refresh_submissions", 200) or 200)),
recent_terminal_window=max(1, int(getattr(args, "gpu_strategy_recent_window", 1000) or 1000)),
long_term_min_samples=max(1, int(getattr(args, "gpu_strategy_min_long_samples", 100) or 100)),
)
strategy_manager.prepare(modelhub_client, supported_gpus=target_gpus, now=now)
strategy_summary = strategy_manager.summary()
if getattr(args, "skip_history_archive", False):
archived_history = []
else:
@@ -561,8 +613,6 @@ def run_submission(
skipped.extend(model_skipped)
failed.extend(model_failed)
write_jsonl(run_dir / "candidates.jsonl", candidates)
submitted: list[dict[str, Any]] = []
duplicate_candidates: list[dict[str, Any]] = []
target_submit_count = resolve_max_submit_count(
@@ -570,8 +620,13 @@ def run_submission(
planned_count=len(candidates),
remaining_daily_quota=remaining_daily_quota,
)
if strategy_manager is not None:
target_submit_count = min(target_submit_count, strategy_manager.submissions_until_refresh)
instance_id = runtime_instance_id()
diversified_candidates = diversify_candidates(candidates, instance_id=instance_id)
if strategy_manager is not None:
diversified_candidates = strategy_manager.order_candidates(diversified_candidates)
write_jsonl(run_dir / "candidates.jsonl", diversified_candidates)
claim_store: SubmissionClaimStore | None = None
attempted_candidates: list[dict[str, Any]] = []
submit_workers = 1
@@ -698,6 +753,8 @@ def run_submission(
claim_store.mark_submitted([*batch_submitted_candidates, *batch_duplicate_candidates])
claim_store.release(batch_failed_candidates)
if strategy_manager is not None:
strategy_manager.record_accepted(batch_submitted_candidates)
if hasattr(modelhub_client, "available_submit_slots") and modelhub_client.available_submit_slots() <= 0:
break
@@ -706,6 +763,9 @@ def run_submission(
write_jsonl(run_dir / "skipped.jsonl", skipped)
write_jsonl(run_dir / "failed.jsonl", failed)
if strategy_manager is not None:
strategy_summary = strategy_manager.summary()
outcome_tracker.save()
summary = {
@@ -723,6 +783,7 @@ def run_submission(
"historyStatsThreshold": getattr(args, "history_stats_threshold", 500),
"historyArchivePath": str(history_archive_path),
"historyArchiveRecordCount": len(archived_history),
"gpuStrategy": strategy_summary,
"scanLimit": scan_limit,
"scannedModels": len(models),
"candidateCount": len(candidates),