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_platform_failures_use_short_circuit_without_poisoning_compatibility_rate(self) -> None: manager = MarketIntelligenceManager("unused.json", framework_min_samples=100) manager.state = { "frameworkStats": { "text-generation": { "gpu": { "vllm": {"modelCount": 1000, "wilsonLowerBound": 0.30}, } } } } manager.set_local_outcome_stats( { "combinationStats": { "gpu|vllm|text-generation": { "successCount": 10, "failureCount": 3, "attributableFailureCount": 0, } }, "recentCombinationStats": { "gpu|vllm|text-generation": { "successCount": 10, "failureCount": 3, "attributableFailureCount": 0, "consecutiveFailures": 0, "consecutivePlatformFailures": 3, "lastPlatformFailureAt": datetime.now(timezone.utc).isoformat(), } }, } ) metadata = manager.framework_metadata("text-generation", "gpu", "vllm") self.assertEqual(10, metadata["frameworkLocalSamples"]) self.assertEqual(1.0, metadata["frameworkLocalSuccessRate"]) self.assertTrue(metadata["frameworkCircuitOpen"]) self.assertEqual( "three_consecutive_platform_failures", metadata["frameworkCircuitReason"], ) 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()