feat: add adaptive GPU scheduling
This commit is contained in:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user