feat: add adaptive GPU scheduling
This commit is contained in:
@@ -14,7 +14,7 @@ if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.remove(str(PACKAGE_DIR))
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from common import read_jsonl # noqa: E402
|
||||
from common import read_json, read_jsonl # noqa: E402
|
||||
from main import build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402
|
||||
from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error # noqa: E402
|
||||
from models import HFModelSummary, ModelInspection # noqa: E402
|
||||
@@ -106,6 +106,54 @@ class DuplicateThenSuccessClient:
|
||||
return {"code": 0, "data": {"id": f"task-{self.calls}"}}
|
||||
|
||||
|
||||
class DynamicCapacityClient:
|
||||
def __init__(self, *, active: int, limit: int, token: str = "dynamic-test") -> None:
|
||||
self.active = active
|
||||
self.limit = limit
|
||||
self.token = token
|
||||
|
||||
def count_active_tasks(self, *, max_count: int, **_kwargs) -> int: # noqa: ANN003
|
||||
return min(self.active, max_count)
|
||||
|
||||
def add_task(self, _payload: dict) -> dict:
|
||||
if self.active >= self.limit:
|
||||
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
|
||||
self.active += 1
|
||||
return {"code": 0, "data": {"id": self.active}}
|
||||
|
||||
|
||||
class AutoStrategyClient:
|
||||
def __init__(self, available: int) -> None:
|
||||
self.available = available
|
||||
self.calls = 0
|
||||
|
||||
@staticmethod
|
||||
def begin_cycle() -> None:
|
||||
return None
|
||||
|
||||
def available_submit_slots(self) -> int:
|
||||
return self.available
|
||||
|
||||
@staticmethod
|
||||
def list_tasks(**_kwargs) -> list[dict]: # noqa: ANN003
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003
|
||||
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||||
|
||||
@staticmethod
|
||||
def processed_gpus_for_model(_model_id: str) -> set[str]:
|
||||
return set()
|
||||
|
||||
def add_task(self, _payload: dict) -> dict:
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise ModelHubAPIError("模型正在验证中,请勿重复提交")
|
||||
self.available -= 1
|
||||
return {"code": 0, "data": {"id": f"strategy-{self.calls}"}}
|
||||
|
||||
|
||||
def make_candidate(index: int) -> dict:
|
||||
return {
|
||||
"repoId": f"owner/model-{index}",
|
||||
@@ -115,6 +163,42 @@ def make_candidate(index: int) -> dict:
|
||||
|
||||
|
||||
class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
def test_adaptive_strategy_counts_only_platform_accepted_tasks(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
root = Path(temporary_dir)
|
||||
args = build_parser().parse_args(
|
||||
[
|
||||
"--task-types",
|
||||
"text-generation",
|
||||
"--limit",
|
||||
"4",
|
||||
"--max-scan-models",
|
||||
"4",
|
||||
"--skip-outcome-sync",
|
||||
"--skip-history-archive",
|
||||
]
|
||||
)
|
||||
args.runs_dir = str(root / "runs")
|
||||
args.ledger_path = str(root / "ledger.jsonl")
|
||||
args.outcomes_path = str(root / "outcomes.jsonl")
|
||||
args.claims_path = str(root / "claims.jsonl")
|
||||
args.history_archive_path = str(root / "history.jsonl")
|
||||
args.gpu_strategy_state_path = str(root / "strategy.json")
|
||||
client = AutoStrategyClient(available=3)
|
||||
|
||||
summary = run_submission(
|
||||
args,
|
||||
now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
|
||||
hf_discovery=FakeDiscovery(4), # type: ignore[arg-type]
|
||||
modelhub_client=client, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
state = read_json(root / "strategy.json")
|
||||
self.assertEqual(3, summary["submittedCount"])
|
||||
self.assertEqual(1, summary["duplicateCount"])
|
||||
self.assertEqual(3, state["acceptedSinceRefresh"])
|
||||
self.assertEqual(3, sum(state["acceptedByCategory"].values()))
|
||||
|
||||
def test_duplicate_batch_is_replaced_until_available_slots_are_filled(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
root = Path(temporary_dir)
|
||||
@@ -231,6 +315,45 @@ class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||
pool.search_by_model_id("owner/model")
|
||||
self.assertEqual(2, client.search_calls)
|
||||
|
||||
def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None:
|
||||
client = DynamicCapacityClient(active=2, limit=3)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[arg-type]
|
||||
active_task_cap=2,
|
||||
active_counts_ttl=60,
|
||||
capacity_probe_interval_cycles=3,
|
||||
instance_id="capacity-growth-test",
|
||||
)
|
||||
|
||||
pool.configure_capacity_probe(2)
|
||||
self.assertEqual(0, pool.available_submit_slots())
|
||||
pool.configure_capacity_probe(3)
|
||||
self.assertEqual(1, pool.available_submit_slots())
|
||||
pool.add_task({"model": "probe"})
|
||||
|
||||
self.assertEqual([3], pool.account_capacity_limits())
|
||||
self.assertEqual([3], pool.active_task_counts())
|
||||
self.assertEqual(0, pool.available_submit_slots())
|
||||
|
||||
def test_rejected_capacity_probe_keeps_known_limit_and_enters_cooldown(self) -> None:
|
||||
client = DynamicCapacityClient(active=2, limit=2)
|
||||
pool = ModelHubClientPool(
|
||||
[client], # type: ignore[arg-type]
|
||||
active_task_cap=2,
|
||||
active_counts_ttl=60,
|
||||
capacity_probe_interval_cycles=3,
|
||||
capacity_probe_cooldown_cycles=3,
|
||||
instance_id="capacity-rejection-test",
|
||||
)
|
||||
|
||||
pool.configure_capacity_probe(3)
|
||||
self.assertEqual(1, pool.available_submit_slots())
|
||||
with self.assertRaises(ModelHubAPIError):
|
||||
pool.add_task({"model": "probe"})
|
||||
|
||||
self.assertEqual([2], pool.account_capacity_limits())
|
||||
self.assertEqual(0, pool.available_submit_slots())
|
||||
|
||||
|
||||
class ProcessCoordinationTests(unittest.TestCase):
|
||||
def test_duplicate_claims_are_retained_so_next_batch_moves_forward(self) -> None:
|
||||
|
||||
122
tests/test_gpu_strategy.py
Normal file
122
tests/test_gpu_strategy.py
Normal file
@@ -0,0 +1,122 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from collections import Counter
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
||||
if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.remove(str(PACKAGE_DIR))
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from gpu_strategy import ( # noqa: E402
|
||||
ALL_SUPPORTED,
|
||||
LONG_TERM,
|
||||
RECENT,
|
||||
GPUStrategyManager,
|
||||
build_strategy_snapshot,
|
||||
choose_next_category,
|
||||
)
|
||||
|
||||
|
||||
def make_history(gpu: str, success: int, failure: int, *, start: datetime) -> list[dict]:
|
||||
records: list[dict] = []
|
||||
for index in range(success + failure):
|
||||
passed = index < success
|
||||
records.append(
|
||||
{
|
||||
"taskId": f"{gpu}-{index}",
|
||||
"gpuType": gpu,
|
||||
"status": "success",
|
||||
"verifyResult": 1 if passed else -1,
|
||||
"updateTime": (start + timedelta(seconds=index)).isoformat(),
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
class HistoryClient:
|
||||
def __init__(self, tasks: list[dict]) -> None:
|
||||
self.tasks = tasks
|
||||
self.list_calls = 0
|
||||
|
||||
def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003
|
||||
self.list_calls += 1
|
||||
return list(self.tasks)
|
||||
|
||||
|
||||
class GPUStrategyTests(unittest.TestCase):
|
||||
def test_long_term_ranking_rejects_tiny_high_rate_sample(self) -> None:
|
||||
start = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
supported = ["tiny", "best", "second", "third", "bad"]
|
||||
tasks = [
|
||||
*make_history("tiny", 8, 0, start=start),
|
||||
*make_history("best", 40, 160, start=start),
|
||||
*make_history("second", 30, 270, start=start),
|
||||
*make_history("third", 45, 455, start=start),
|
||||
*make_history("bad", 5, 495, start=start),
|
||||
]
|
||||
|
||||
snapshot = build_strategy_snapshot(
|
||||
tasks,
|
||||
supported_gpus=supported,
|
||||
long_term_min_samples=100,
|
||||
)
|
||||
|
||||
self.assertEqual(["best", "second", "third"], snapshot["longTermGpus"])
|
||||
self.assertNotIn("tiny", snapshot["longTermGpus"])
|
||||
|
||||
def test_weighted_category_planner_is_exact_over_200_accepts(self) -> None:
|
||||
counts = {LONG_TERM: 0, ALL_SUPPORTED: 0, RECENT: 0}
|
||||
for _ in range(200):
|
||||
category = choose_next_category(counts)
|
||||
counts[category] += 1
|
||||
|
||||
self.assertEqual({LONG_TERM: 100, ALL_SUPPORTED: 60, RECENT: 40}, counts)
|
||||
|
||||
def test_candidate_order_uses_50_30_20_for_first_strategy_generation(self) -> None:
|
||||
supported = ["gpu-a", "gpu-b", "gpu-c", "gpu-d"]
|
||||
snapshot = build_strategy_snapshot([], supported_gpus=supported)
|
||||
manager = GPUStrategyManager("unused.json")
|
||||
manager.state = snapshot
|
||||
candidates = [
|
||||
{"repoId": f"owner/model-{model}", "targetGpu": gpu}
|
||||
for model in range(250)
|
||||
for gpu in supported
|
||||
]
|
||||
|
||||
ordered = manager.order_candidates(candidates)
|
||||
categories = Counter(candidate["strategyCategory"] for candidate in ordered[:200])
|
||||
|
||||
self.assertEqual(100, categories[LONG_TERM])
|
||||
self.assertEqual(60, categories[ALL_SUPPORTED])
|
||||
self.assertEqual(40, categories[RECENT])
|
||||
|
||||
def test_refresh_happens_only_after_200_accepted_submissions(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||
path = Path(temporary_dir) / "strategy.json"
|
||||
client = HistoryClient([])
|
||||
manager = GPUStrategyManager(path, refresh_submissions=200)
|
||||
manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"])
|
||||
self.assertEqual(1, client.list_calls)
|
||||
|
||||
manager.record_accepted([{"strategyCategory": LONG_TERM}] * 199)
|
||||
self.assertEqual(1, manager.submissions_until_refresh)
|
||||
manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"])
|
||||
self.assertEqual(1, client.list_calls)
|
||||
|
||||
manager.record_accepted([{"strategyCategory": RECENT}])
|
||||
self.assertEqual(0, manager.submissions_until_refresh)
|
||||
manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"])
|
||||
self.assertEqual(2, client.list_calls)
|
||||
self.assertEqual(1, manager.state["generation"])
|
||||
self.assertEqual(200, manager.state["acceptedTotal"])
|
||||
self.assertEqual(0, manager.state["acceptedSinceRefresh"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -11,6 +11,7 @@ if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from hf_discovery import HuggingFaceDiscovery # noqa: E402
|
||||
from http_json import HttpJsonError # noqa: E402
|
||||
|
||||
|
||||
class RecordingHttpClient:
|
||||
@@ -22,6 +23,30 @@ class RecordingHttpClient:
|
||||
return {"success": True, "data": {"models": []}}
|
||||
|
||||
|
||||
class RateLimitedPageClient:
|
||||
def __init__(self) -> None:
|
||||
self.page_calls: list[int] = []
|
||||
self.page_two_attempts = 0
|
||||
|
||||
def request_json(self, _method: str, _path: str, *, query: dict) -> dict:
|
||||
page = int(query["page_number"])
|
||||
self.page_calls.append(page)
|
||||
if page == 2:
|
||||
self.page_two_attempts += 1
|
||||
if self.page_two_attempts == 1:
|
||||
raise HttpJsonError("rate limited", status_code=429)
|
||||
start = (page - 1) * 50
|
||||
items = [
|
||||
{
|
||||
"id": f"owner/model-{index}",
|
||||
"downloads": 100,
|
||||
"last_modified": "2026-01-01T00:00:00Z",
|
||||
}
|
||||
for index in range(start, start + 50)
|
||||
]
|
||||
return {"success": True, "data": {"models": items}}
|
||||
|
||||
|
||||
class ModelScopeDiscoveryTests(unittest.TestCase):
|
||||
def test_openapi_page_size_never_exceeds_platform_limit(self) -> None:
|
||||
http_client = RecordingHttpClient()
|
||||
@@ -36,6 +61,29 @@ class ModelScopeDiscoveryTests(unittest.TestCase):
|
||||
self.assertEqual([], models)
|
||||
self.assertEqual(50, http_client.queries[0]["page_size"])
|
||||
|
||||
def test_rate_limited_page_keeps_previous_pages_and_resumes_from_cache(self) -> None:
|
||||
http_client = RateLimitedPageClient()
|
||||
discovery = HuggingFaceDiscovery(
|
||||
http_client=http_client, # type: ignore[arg-type]
|
||||
page_interval_seconds=0,
|
||||
page_cache_ttl_seconds=60,
|
||||
)
|
||||
|
||||
partial = discovery.list_recent_models(
|
||||
pipeline_tags=["text-generation"],
|
||||
limit=100,
|
||||
min_downloads=0,
|
||||
)
|
||||
resumed = discovery.list_recent_models(
|
||||
pipeline_tags=["text-generation"],
|
||||
limit=100,
|
||||
min_downloads=0,
|
||||
)
|
||||
|
||||
self.assertEqual(50, len(partial))
|
||||
self.assertEqual(100, len(resumed))
|
||||
self.assertEqual([1, 2, 2], http_client.page_calls)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user