feat: prioritize proven GPU framework combinations

This commit is contained in:
CoolBoy
2026-08-05 18:21:59 +08:00
parent 882479e43e
commit 5e47d9e695
11 changed files with 431 additions and 86 deletions

View File

@@ -14,7 +14,6 @@ if str(PACKAGE_DIR) in sys.path:
sys.path.insert(0, str(PACKAGE_DIR))
from gpu_strategy import ( # noqa: E402
ALL_SUPPORTED,
LONG_TERM,
RECENT,
GPUStrategyManager,
@@ -55,7 +54,15 @@ class WeightedMarket:
return {"gpu-a": 3.0, "gpu-b": 0.1, "gpu-c": 0.1, "gpu-d": 0.1}[gpu]
def gpu_metadata(self, gpu: str) -> dict:
return {"marketWeight": self.gpu_weight(gpu, category=ALL_SUPPORTED)}
return {"marketWeight": self.gpu_weight(gpu, category=LONG_TERM)}
def eligible_gpus(self, supported_gpus: list[str]) -> list[str]:
return list(supported_gpus)
class VettedMarket(WeightedMarket):
def eligible_gpus(self, supported_gpus: list[str]) -> list[str]:
return [gpu for gpu in supported_gpus if gpu in {"gpu-b", "gpu-c"}]
class GPUStrategyTests(unittest.TestCase):
@@ -80,14 +87,14 @@ class GPUStrategyTests(unittest.TestCase):
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}
counts = {LONG_TERM: 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)
self.assertEqual({LONG_TERM: 140, RECENT: 60}, counts)
def test_candidate_order_uses_50_30_20_for_first_strategy_generation(self) -> None:
def test_candidate_order_uses_70_30_without_exploration(self) -> None:
supported = ["gpu-a", "gpu-b", "gpu-c", "gpu-d"]
snapshot = build_strategy_snapshot([], supported_gpus=supported)
manager = GPUStrategyManager("unused.json")
@@ -101,9 +108,9 @@ class GPUStrategyTests(unittest.TestCase):
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])
self.assertEqual(140, categories[LONG_TERM])
self.assertEqual(60, categories[RECENT])
self.assertNotIn("all_supported", categories)
def test_refresh_happens_only_after_200_accepted_submissions(self) -> None:
with tempfile.TemporaryDirectory() as temporary_dir:
@@ -126,7 +133,7 @@ class GPUStrategyTests(unittest.TestCase):
self.assertEqual(200, manager.state["acceptedTotal"])
self.assertEqual(0, manager.state["acceptedSinceRefresh"])
def test_market_weights_change_gpu_mix_without_changing_50_30_20_split(self) -> None:
def test_market_weights_change_gpu_mix_without_changing_70_30_split(self) -> None:
supported = ["gpu-a", "gpu-b", "gpu-c", "gpu-d"]
manager = GPUStrategyManager("unused.json", market_intelligence=WeightedMarket())
manager.state = build_strategy_snapshot([], supported_gpus=supported)
@@ -138,15 +145,26 @@ class GPUStrategyTests(unittest.TestCase):
ordered = manager.order_candidates(candidates)[:200]
category_counts = Counter(candidate["strategyCategory"] for candidate in ordered)
exploration_counts = Counter(
candidate["targetGpu"]
for candidate in ordered
if candidate["strategyCategory"] == ALL_SUPPORTED
)
gpu_counts = Counter(candidate["targetGpu"] for candidate in ordered)
self.assertEqual({LONG_TERM: 100, ALL_SUPPORTED: 60, RECENT: 40}, category_counts)
self.assertGreater(exploration_counts["gpu-a"], exploration_counts["gpu-b"])
self.assertTrue(all(exploration_counts[gpu] > 0 for gpu in supported))
self.assertEqual({LONG_TERM: 140, RECENT: 60}, category_counts)
self.assertGreater(gpu_counts["gpu-a"], gpu_counts["gpu-b"])
self.assertNotIn("all_supported", category_counts)
def test_unvetted_gpus_receive_zero_submissions(self) -> None:
supported = ["gpu-a", "gpu-b", "gpu-c", "gpu-d"]
manager = GPUStrategyManager("unused.json", market_intelligence=VettedMarket())
manager.state = build_strategy_snapshot([], supported_gpus=supported)
candidates = [
{"repoId": f"owner/model-{model}", "targetGpu": gpu}
for model in range(50)
for gpu in supported
]
ordered = manager.order_candidates(candidates)
self.assertTrue(ordered)
self.assertEqual({"gpu-b", "gpu-c"}, {candidate["targetGpu"] for candidate in ordered})
if __name__ == "__main__":

View File

@@ -13,6 +13,7 @@ if str(PACKAGE_DIR) in sys.path:
sys.path.insert(0, str(PACKAGE_DIR))
from main import choose_candidate_for_gpu # noqa: E402
from common import write_json # noqa: E402
from market_intelligence import MarketIntelligenceManager # noqa: E402
from models import HFModelSummary, ModelInspection # noqa: E402
from template_selector import TemplateSelector # noqa: E402
@@ -158,6 +159,52 @@ class MarketIntelligenceTests(unittest.TestCase):
)
self.assertEqual(first_call_count, client.calls)
def test_version_two_snapshot_is_migrated_without_losing_public_evidence(self) -> None:
with tempfile.TemporaryDirectory() as temporary_dir:
path = Path(temporary_dir) / "market.json"
now = datetime.now(timezone.utc)
write_json(
path,
{
"version": 2,
"supportedGpus": ["fast"],
"taskTypes": ["text-generation"],
"queueUpdatedAt": now.isoformat(),
"frameworkUpdatedAt": now.isoformat(),
"queueError": None,
"frameworkError": None,
"gpuStats": {
"fast": {
"gpu": "fast",
"canVerify": True,
"recentTerminal": 100,
"recentSuccess": 80,
"recentWilsonLowerBound": 0.70,
"backlogHours": 10.0,
"healthFactor": 1.0,
}
},
"frameworkStats": {
"text-generation": {
"fast": {
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.20}
}
}
},
},
)
client = FailingMarketClient()
manager = MarketIntelligenceManager(path, log_fn=lambda _message: None)
state = manager.prepare(
client,
supported_gpus=["fast"],
task_types=["text-generation"],
now=now + timedelta(seconds=30),
)
self.assertEqual(3, state["version"])
self.assertTrue(state["gpuStats"]["fast"]["submissionEligible"])
self.assertEqual(0, client.calls)
def test_framework_ranking_uses_confidence_bound_and_ignores_tiny_samples(self) -> None:
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
manager.state = {
@@ -332,6 +379,76 @@ class MarketIntelligenceTests(unittest.TestCase):
assert candidate is not None
self.assertEqual("vllm", candidate.framework)
def test_unproven_public_framework_receives_no_submission(self) -> None:
manager = MarketIntelligenceManager("unused.json")
manager.state = {
"frameworkStats": {
"text-generation": {
"Cambricon_mlu-370-x4": {
"vllm": {"modelCount": 299, "wilsonLowerBound": 0.50},
}
}
}
}
candidate = choose_candidate_for_gpu(
model=HFModelSummary(
repo_id="owner/model",
downloads=100,
last_modified=None,
pipeline_tag="text-generation",
),
inspection=ModelInspection(repo_id="owner/model", weight_files=["model.safetensors"]),
template_selector=TemplateSelector(),
task_types=["text-generation"],
target_gpu="Cambricon_mlu-370-x4",
market_intelligence=manager,
)
self.assertIsNone(candidate)
def test_five_consecutive_local_failures_circuit_break_framework(self) -> None:
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
manager.state = {
"frameworkStats": {
"text-generation": {
"Iluvatar_bi-100": {
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.40},
"transformers": {"modelCount": 1000, "wilsonLowerBound": 0.20},
}
}
}
}
manager.set_local_outcome_stats(
{
"combinationStats": {},
"recentCombinationStats": {
"Iluvatar_bi-100|vllm|text-generation": {
"successCount": 0,
"failureCount": 5,
"consecutiveFailures": 5,
"lastTerminalAt": datetime.now(timezone.utc).isoformat(),
}
},
}
)
candidate = choose_candidate_for_gpu(
model=HFModelSummary(
repo_id="owner/model",
downloads=100,
last_modified=None,
pipeline_tag="text-generation",
),
inspection=ModelInspection(repo_id="owner/model", weight_files=["model.safetensors"]),
template_selector=TemplateSelector(),
task_types=["text-generation"],
target_gpu="Iluvatar_bi-100",
market_intelligence=manager,
)
self.assertIsNotNone(candidate)
assert candidate is not None
self.assertEqual("transformers", candidate.framework)
metadata = manager.framework_metadata("text-generation", "Iluvatar_bi-100", "vllm")
self.assertTrue(metadata["frameworkCircuitOpen"])
if __name__ == "__main__":
unittest.main()