feat: add adaptive GPU scheduling

modelhub_submmit_api/poll_runner.py

@@ -10,9 +10,10 @@ from typing import Any, Callable
 
 from common import utc_now, write_json
 from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
+from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH
 from hf_discovery import HuggingFaceDiscovery
 from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir
-from modelhub_client import ModelHubClient, ModelHubClientPool
+from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool
 from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
 from runner_common import DEFAULT_KEY_PATH, ensure_tokens
 from submission_claims import DEFAULT_CLAIMS_PATH
@@ -58,6 +59,13 @@ def build_parser() -> argparse.ArgumentParser:
     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("--disable-gpu-strategy", action="store_true", help="Disable adaptive 50/30/20 GPU scheduling")
+    parser.add_argument(
+        "--gpu-strategy-refresh-submissions",
+        type=int,
+        default=200,
+        help="Recalculate adaptive GPU choices after this many accepted submissions",
+    )
     parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_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)
@@ -68,6 +76,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("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
     parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS)
     parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
@@ -97,7 +123,11 @@ def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClientPool:
     modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
     token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token]
     clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values]
-    return ModelHubClientPool(clients)
+    return ModelHubClientPool(
+        clients,
+        capacity_probe_interval_cycles=max(0, int(getattr(base_args, "capacity_probe_interval_cycles", 3) or 0)),
+        capacity_state_path=Path(getattr(base_args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)),
+    )
 
 
 def run_poll_loop(
@@ -140,10 +170,16 @@ def run_poll_loop(
             break
 
         cycles += 1
+        if hasattr(modelhub_client, "configure_capacity_probe"):
+            modelhub_client.configure_capacity_probe(cycles)
         active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else []
+        capacity_limits = modelhub_client.account_capacity_limits() if hasattr(modelhub_client, "account_capacity_limits") else []
+        capacity_probe = modelhub_client.capacity_probe_enabled() if hasattr(modelhub_client, "capacity_probe_enabled") else False
         available_slots = modelhub_client.available_submit_slots() if hasattr(modelhub_client, "available_submit_slots") else None
         log(
             f"[poll] cycle={cycles} active_counts={','.join(str(count) for count in active_counts) if active_counts else 'n/a'} "
+            f"capacity_limits={','.join(str(limit) for limit in capacity_limits) if capacity_limits else 'n/a'} "
+            f"capacity_probe={'on' if capacity_probe else 'off'} "
             f"available_slots={available_slots if available_slots is not None else 'n/a'}"
         )
 
@@ -152,8 +188,10 @@ def run_poll_loop(
             time.sleep(base_args.poll_interval_seconds)
             continue
 
+        cycle_args = _make_cycle_args(base_args)
+        cycle_args.capacity_probe_already_configured = True
         cycle_summary = run_fn(
-            base_args=_make_cycle_args(base_args),
+            base_args=cycle_args,
             now=utc_now(),
             hf_discovery=hf_discovery,
             modelhub_client=modelhub_client,