455 lines
18 KiB
Python
455 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
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 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
|
|
|
|
|
|
class PublicMarketClient:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
self.counts = {
|
|
"fast": {"waiting": 100, "running": 8, "completed": 120, "success": 80, "failed": 0},
|
|
"slow": {"waiting": 10, "running": 1, "completed": 1, "success": 1, "failed": 0},
|
|
"stopped": {"waiting": 500, "running": 0, "completed": 0, "success": 0, "failed": 0},
|
|
}
|
|
|
|
def list_machine_info(self) -> list[dict]:
|
|
self.calls += 1
|
|
return [
|
|
{"gpuType": "fast", "canVerify": True, "maxConcurrentTasks": 8},
|
|
{"gpuType": "slow", "canVerify": True, "maxConcurrentTasks": 1},
|
|
{"gpuType": "stopped", "canVerify": False, "maxConcurrentTasks": 0},
|
|
]
|
|
|
|
def list_tasks_page(self, **kwargs) -> dict: # noqa: ANN003
|
|
self.calls += 1
|
|
row = self.counts[kwargs["gpu_type"]]
|
|
status = kwargs.get("status")
|
|
verify_result = kwargs.get("verify_result")
|
|
if status == "waiting":
|
|
total = row["waiting"]
|
|
elif status == "running":
|
|
total = row["running"]
|
|
elif status == "failed":
|
|
total = row["failed"]
|
|
elif status == "success" and verify_result == 1:
|
|
total = row["success"]
|
|
else:
|
|
total = row["completed"]
|
|
return {"data": {"total": total, "records": []}}
|
|
|
|
def list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict]:
|
|
self.calls += 1
|
|
del task_type, target_gpu
|
|
return [
|
|
{"framework": "vllm", "modelCount": 1000, "successCount": 100},
|
|
{"framework": "transformers", "modelCount": 1000, "successCount": 300},
|
|
]
|
|
|
|
def get_build_config(self, task_type: str, target_gpu: str, framework: str) -> str:
|
|
self.calls += 1
|
|
del task_type, target_gpu
|
|
return (
|
|
f"framework: {framework}\n"
|
|
"sut_config:\n gpu_num: 1\n values: {}\n"
|
|
"ref_config:\n gpu_num: 1\n values: {}\n"
|
|
)
|
|
|
|
|
|
class FailingMarketClient:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def list_machine_info(self) -> list[dict]:
|
|
self.calls += 1
|
|
raise RuntimeError("temporary queue outage")
|
|
|
|
def list_tasks_page(self, **_kwargs) -> dict: # noqa: ANN003
|
|
self.calls += 1
|
|
raise RuntimeError("temporary queue outage")
|
|
|
|
def list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict]:
|
|
self.calls += 1
|
|
del task_type, target_gpu
|
|
raise RuntimeError("temporary framework outage")
|
|
|
|
|
|
class MarketIntelligenceTests(unittest.TestCase):
|
|
def test_expected_throughput_beats_short_raw_queue_and_stopped_gpu_is_demoted(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
client = PublicMarketClient()
|
|
manager = MarketIntelligenceManager(
|
|
Path(temporary_dir) / "market.json",
|
|
throughput_window_hours=6,
|
|
log_fn=lambda _message: None,
|
|
)
|
|
state = manager.prepare(
|
|
client,
|
|
supported_gpus=["fast", "slow", "stopped"],
|
|
task_types=["text-generation"],
|
|
now=datetime(2026, 8, 4, tzinfo=timezone.utc),
|
|
)
|
|
|
|
fast = state["gpuStats"]["fast"]
|
|
slow = state["gpuStats"]["slow"]
|
|
stopped = state["gpuStats"]["stopped"]
|
|
self.assertGreater(fast["throughputPerHour"], slow["throughputPerHour"])
|
|
self.assertLess(fast["backlogHours"], slow["backlogHours"])
|
|
self.assertGreater(fast["selectionWeight"], slow["selectionWeight"])
|
|
self.assertLessEqual(stopped["selectionWeight"], 0.1)
|
|
self.assertTrue(
|
|
state["frameworkStats"]["text-generation"]["fast"]["transformers"]["officialConfigValid"]
|
|
)
|
|
|
|
def test_snapshot_cache_prevents_repeated_public_api_scans(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
client = PublicMarketClient()
|
|
manager = MarketIntelligenceManager(Path(temporary_dir) / "market.json", log_fn=lambda _message: None)
|
|
started = datetime(2026, 8, 4, tzinfo=timezone.utc)
|
|
manager.prepare(
|
|
client,
|
|
supported_gpus=["fast", "slow"],
|
|
task_types=["text-generation"],
|
|
now=started,
|
|
)
|
|
first_call_count = client.calls
|
|
manager.prepare(
|
|
client,
|
|
supported_gpus=["fast", "slow"],
|
|
task_types=["text-generation"],
|
|
now=started + timedelta(seconds=30),
|
|
)
|
|
self.assertEqual(first_call_count, client.calls)
|
|
|
|
def test_market_outage_falls_back_to_neutral_and_uses_retry_backoff(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
|
client = FailingMarketClient()
|
|
manager = MarketIntelligenceManager(Path(temporary_dir) / "market.json", log_fn=lambda _message: None)
|
|
started = datetime(2026, 8, 4, tzinfo=timezone.utc)
|
|
state = manager.prepare(
|
|
client,
|
|
supported_gpus=["gpu"],
|
|
task_types=["text-generation"],
|
|
now=started,
|
|
)
|
|
self.assertEqual(1.0, state["gpuStats"]["gpu"]["selectionWeight"])
|
|
self.assertIsNotNone(state["queueError"])
|
|
self.assertIsNotNone(state["frameworkError"])
|
|
first_call_count = client.calls
|
|
manager.prepare(
|
|
client,
|
|
supported_gpus=["gpu"],
|
|
task_types=["text-generation"],
|
|
now=started + timedelta(seconds=30),
|
|
)
|
|
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 = {
|
|
"frameworkStats": {
|
|
"text-generation": {
|
|
"gpu": {
|
|
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.10},
|
|
"vllm-mlu": {"modelCount": 800, "wilsonLowerBound": 0.35},
|
|
"vllm-customized": {"modelCount": 2, "wilsonLowerBound": 0.90},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
ranked = manager.rank_frameworks(
|
|
task_type="text-generation",
|
|
target_gpu="gpu",
|
|
compatible_frameworks=["vllm", "vllm-customized", "vllm-mlu"],
|
|
)
|
|
self.assertEqual(["vllm-mlu", "vllm", "vllm-customized"], ranked)
|
|
|
|
def test_local_account_evidence_is_blended_without_overriding_sample_guards(self) -> None:
|
|
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
|
|
manager.state = {
|
|
"frameworkStats": {
|
|
"text-generation": {
|
|
"gpu": {
|
|
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.30},
|
|
"vllm-mlu": {"modelCount": 1000, "wilsonLowerBound": 0.35},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
manager.set_local_outcome_stats(
|
|
{
|
|
"combinationStats": {
|
|
"gpu|vllm|text-generation": {
|
|
"successCount": 20,
|
|
"failureCount": 0,
|
|
}
|
|
}
|
|
}
|
|
)
|
|
ranked = manager.rank_frameworks(
|
|
task_type="text-generation",
|
|
target_gpu="gpu",
|
|
compatible_frameworks=["vllm", "vllm-mlu"],
|
|
)
|
|
self.assertEqual("vllm", ranked[0])
|
|
metadata = manager.framework_metadata("text-generation", "gpu", "vllm")
|
|
self.assertEqual(20, metadata["frameworkLocalSamples"])
|
|
self.assertGreater(metadata["frameworkCombinedScore"], 0.35)
|
|
|
|
def test_candidate_uses_best_supported_public_framework(self) -> None:
|
|
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
|
|
manager.state = {
|
|
"frameworkStats": {
|
|
"text-generation": {
|
|
"Cambricon_mlu-370-x4": {
|
|
"vllm": {
|
|
"modelCount": 1000,
|
|
"successCount": 100,
|
|
"successRate": 0.10,
|
|
"wilsonLowerBound": 0.08,
|
|
},
|
|
"vllm-mlu": {
|
|
"modelCount": 1000,
|
|
"successCount": 400,
|
|
"successRate": 0.40,
|
|
"wilsonLowerBound": 0.37,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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.assertIsNotNone(candidate)
|
|
assert candidate is not None
|
|
self.assertEqual("vllm-mlu", candidate.framework)
|
|
self.assertAlmostEqual(0.37, candidate.score)
|
|
|
|
def test_new_framework_is_discovered_but_only_wins_on_qualified_success_score(self) -> None:
|
|
official_config = (
|
|
"framework: future-engine\n"
|
|
"sut_config:\n gpu_num: 1\n values: {}\n"
|
|
"ref_config:\n gpu_num: 1\n values: {}\n"
|
|
)
|
|
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
|
|
manager.state = {
|
|
"frameworkStats": {
|
|
"text-generation": {
|
|
"Cambricon_mlu-370-x4": {
|
|
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.10},
|
|
"future-engine": {
|
|
"modelCount": 1000,
|
|
"successCount": 500,
|
|
"successRate": 0.50,
|
|
"wilsonLowerBound": 0.47,
|
|
"officialConfigValid": True,
|
|
"officialConfig": official_config,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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.assertIsNotNone(candidate)
|
|
assert candidate is not None
|
|
self.assertEqual("future-engine", candidate.framework)
|
|
self.assertEqual(official_config, candidate.config_params)
|
|
self.assertIn("official_build_config_synced", candidate.warnings)
|
|
|
|
def test_tiny_new_framework_sample_does_not_displace_safe_legacy_framework(self) -> None:
|
|
manager = MarketIntelligenceManager("unused.json", framework_min_samples=100)
|
|
manager.state = {
|
|
"frameworkStats": {
|
|
"text-generation": {
|
|
"Cambricon_mlu-370-x4": {
|
|
"vllm": {"modelCount": 1000, "wilsonLowerBound": 0.10},
|
|
"future-engine": {
|
|
"modelCount": 2,
|
|
"wilsonLowerBound": 0.90,
|
|
"officialConfigValid": True,
|
|
"officialConfig": (
|
|
"framework: future-engine\n"
|
|
"sut_config:\n gpu_num: 1\n values: {}\n"
|
|
"ref_config:\n gpu_num: 1\n values: {}\n"
|
|
),
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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.assertIsNotNone(candidate)
|
|
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()
|