from __future__ import annotations import sys import tempfile import unittest from datetime import datetime, 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 build_parser, one_candidate_per_model, process_model_for_candidates, run_submission, submit_candidate # noqa: E402 from candidate_preflight import CandidatePreflightAdvisor # noqa: E402 from modelhub_client import ( # noqa: E402 ModelHubAPIError, ModelHubClientPool, is_model_uniqueness_error, parse_model_submission_precheck, ) from models import HFModelSummary, ModelInspection # noqa: E402 from submission_exclusions import SubmissionExclusionStore # noqa: E402 from template_selector import TemplateSelector # noqa: E402 class SafeDiscovery: def __init__(self) -> None: self.model = HFModelSummary( repo_id="owner/model", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) def list_recent_models(self, **_kwargs) -> list[HFModelSummary]: # noqa: ANN003 return [self.model] @staticmethod def inspect_model(model: HFModelSummary) -> ModelInspection: return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"]) class OtherGpuProcessedClient: @staticmethod def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict: del force_refresh return {"processedGpus": {"Biren_166m"}, "isInDB": True} class FailingLookupClient: @staticmethod def search_by_model_id(_model_id: str) -> dict: raise ModelHubAPIError("temporary lookup outage") class ExactGpuProcessedClient: def __init__(self) -> None: self.add_calls = 0 @staticmethod def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict: del force_refresh return {"processedGpus": {"Vastai_va16"}, "isInDB": True} def add_task(self, _payload: dict) -> dict: self.add_calls += 1 return {"code": 0, "data": {"id": "should-not-submit"}} class FailedFrameworkClient: def __init__(self) -> None: self.add_calls = 0 @staticmethod def model_submission_precheck(_model_id: str, *, force_refresh: bool = False) -> dict: del force_refresh return { "processedGpus": set(), "failedFrameworksByGpu": {"Cambricon_mlu-370-x4": {"vllm"}}, } def add_task(self, _payload: dict) -> dict: self.add_calls += 1 return {"code": 0, "data": {"id": str(self.add_calls)}} class AllGpuOfficialRegistry: @staticmethod def task_types_for(_client, *, model_address: str, model_last_modified: str | None, gpu: str) -> list[str]: del model_address, model_last_modified, gpu return ["text-generation"] class AllGpuVettedMarket: @staticmethod def has_framework_route(_task_type: str, _target_gpu: str) -> bool: return True @staticmethod def ensure_framework_route(_client, _task_type: str, _target_gpu: str) -> None: return None @staticmethod def selectable_frameworks( *, task_type: str, target_gpu: str, incumbent_frameworks: list[str], inspection ) -> list[str]: del task_type, target_gpu, incumbent_frameworks, inspection return ["vllm"] @staticmethod def official_config( *, task_type: str, target_gpu: str, framework: str, gguf_filename: str | None = None ) -> str: del task_type, target_gpu, framework, gguf_filename return "framework: vllm" @staticmethod def framework_metadata(_task_type: str, _target_gpu: str, _framework: str) -> dict: return { "frameworkCombinedScore": 1.0, "frameworkOfficialConfigValid": True, "frameworkEvidenceQualified": True, } @staticmethod def gpu_metadata(_target_gpu: str) -> dict: return {} class SeventhGpuFallbackMarket(AllGpuVettedMarket): @staticmethod def has_framework_route(_task_type: str, target_gpu: str) -> bool: return target_gpu == "gpu-7" class TwoFrameworkMarket(AllGpuVettedMarket): @staticmethod def selectable_frameworks( *, task_type: str, target_gpu: str, incumbent_frameworks: list[str], inspection ) -> list[str]: del task_type, target_gpu, incumbent_frameworks, inspection return ["vllm", "vllm-mlu"] class UniquenessRejectingClient: @staticmethod def add_task(_payload: dict) -> dict: raise ModelHubAPIError("模型唯一性检查没有通过，无法进行同步") class UniquenessRunClient(UniquenessRejectingClient): def __init__(self) -> None: self.add_calls = 0 @staticmethod def available_submit_slots() -> int: return 1 @staticmethod def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003 return {"code": 0, "data": {"records": []}} @staticmethod def processed_gpus_for_model(_model_id: str) -> set[str]: return set() def add_task(self, _payload: dict) -> dict: self.add_calls += 1 return super().add_task(_payload) def candidate(target_gpu: str = "Vastai_va16") -> dict: return { "repoId": "owner/model", "modelAddress": "https://modelscope.cn/models/owner/model", "taskType": "text-generation", "targetGpu": target_gpu, "framework": "vllm", "configParams": "framework: vllm", } class SubmissionSafetyTests(unittest.TestCase): def test_official_routes_prefer_first_six_gpus(self) -> None: target_gpus = [f"gpu-{index}" for index in range(1, 8)] model = HFModelSummary( repo_id="owner/model", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) candidates, skipped, failed = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=OtherGpuProcessedClient(), # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=target_gpus, allowed_task_types=["text-generation"], market_intelligence=AllGpuVettedMarket(), # type: ignore[arg-type] official_registry=AllGpuOfficialRegistry(), # type: ignore[arg-type] ) self.assertEqual(target_gpus[:6], [item["targetGpu"] for item in candidates]) self.assertEqual([], skipped) self.assertEqual([], failed) def test_official_routes_fall_back_after_preferred_six_are_unusable(self) -> None: target_gpus = [f"gpu-{index}" for index in range(1, 8)] model = HFModelSummary( repo_id="owner/model", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) candidates, skipped, failed = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=OtherGpuProcessedClient(), # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=target_gpus, allowed_task_types=["text-generation"], market_intelligence=SeventhGpuFallbackMarket(), # type: ignore[arg-type] official_registry=AllGpuOfficialRegistry(), # type: ignore[arg-type] ) self.assertEqual(["gpu-7"], [item["targetGpu"] for item in candidates]) self.assertEqual(6, len(skipped)) self.assertEqual([], failed) def test_precheck_tracks_processed_gpus_without_blocking_other_gpus(self) -> None: payload = { "code": 0, "data": { "isInDB": True, "verifyResult": { "Biren_166m": { "result": "已验证", "records": [{"verifyResult": 1}], } }, }, } precheck = parse_model_submission_precheck(payload) self.assertEqual({"Biren_166m"}, precheck["processedGpus"]) model = HFModelSummary( repo_id="owner/model", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) candidates, skipped, failed = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=OtherGpuProcessedClient(), # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["Biren_166m", "Vastai_va16"], allowed_task_types=["text-generation"], ) self.assertEqual(["Vastai_va16"], [item["targetGpu"] for item in candidates]) self.assertEqual("already_processed_for_gpu", skipped[0]["reason"]) self.assertEqual([], failed) def test_precheck_distinguishes_unverified_failed_and_active_gpus(self) -> None: payload = { "data": { "isInDB": True, "verifyResult": { "free": {"result": "未验证", "records": []}, "failed": { "result": "验证失败", "records": [{"framework": "vllm", "verifyResult": -1}], }, "running": {"result": "验证中", "records": []}, "success": { "result": "已验证", "records": [{"framework": "vllm", "verifyResult": 1}], }, "unknown": {"result": "未知状态", "records": []}, "failed_without_framework": { "result": "验证失败", "records": [{"verifyResult": -1}], }, }, }, } precheck = parse_model_submission_precheck(payload) self.assertEqual( {"running", "success", "unknown", "failed_without_framework"}, precheck["processedGpus"], ) self.assertEqual({"failed": {"vllm"}}, precheck["failedFrameworksByGpu"]) def test_unverified_gpu_becomes_candidate_but_validating_gpu_does_not(self) -> None: class StatusClient: def __init__(self, status: str) -> None: self.status = status def model_submission_precheck(self, _model_id: str) -> dict: return parse_model_submission_precheck({ "data": { "verifyResult": { "Vastai_va16": {"result": self.status, "records": []} } } }) def collect(status: str): return process_model_for_candidates( model=SafeDiscovery().model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=StatusClient(status), # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["Vastai_va16"], allowed_task_types=["text-generation"], ) open_candidates, _, _ = collect("未验证") active_candidates, active_skips, _ = collect("验证中") self.assertEqual(["Vastai_va16"], [item["targetGpu"] for item in open_candidates]) self.assertEqual([], active_candidates) self.assertEqual("already_processed_for_gpu", active_skips[0]["reason"]) def test_failed_gpu_requires_an_alternative_framework(self) -> None: model = SafeDiscovery().model client = FailedFrameworkClient() candidates, skipped, failed = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["Cambricon_mlu-370-x4"], allowed_task_types=["text-generation"], market_intelligence=TwoFrameworkMarket(), # type: ignore[arg-type] official_registry=AllGpuOfficialRegistry(), # type: ignore[arg-type] ) self.assertEqual([], skipped) self.assertEqual([], failed) self.assertEqual(["vllm-mlu"], [item["framework"] for item in candidates]) self.assertEqual(["vllm"], candidates[0]["previousFailedFrameworks"]) blocked = submit_candidate({**candidates[0], "framework": "vllm"}, client) # type: ignore[arg-type] accepted = submit_candidate(candidates[0], client) # type: ignore[arg-type] self.assertEqual("duplicate", blocked["outcome"]) self.assertEqual("prior_failed_framework_for_gpu", blocked["reason"]) self.assertEqual("submitted", accepted["outcome"]) self.assertEqual(1, client.add_calls) unvetted = submit_candidate( {**candidates[0], "frameworkEvidenceQualified": False}, client # type: ignore[arg-type] ) self.assertEqual("framework_evidence_deferred", unvetted["outcome"]) self.assertEqual(1, client.add_calls) no_routes, no_route_skips, _ = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["Cambricon_mlu-370-x4"], allowed_task_types=["text-generation"], market_intelligence=AllGpuVettedMarket(), # type: ignore[arg-type] official_registry=AllGpuOfficialRegistry(), # type: ignore[arg-type] ) self.assertEqual([], no_routes) self.assertEqual("prior_failed_framework_no_alternative", no_route_skips[-1]["reason"]) no_market_routes, no_market_skips, _ = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["Cambricon_mlu-370-x4"], allowed_task_types=["text-generation"], ) self.assertEqual([], no_market_routes) self.assertEqual( "prior_failed_gpu_requires_public_framework_evidence", no_market_skips[-1]["reason"], ) def test_published_size_avoids_official_calls_for_certain_oom(self) -> None: class CountingRegistry(AllGpuOfficialRegistry): calls = 0 def task_types_for(self, *_args, **_kwargs) -> list[str]: self.calls += 1 return ["text-generation"] registry = CountingRegistry() model = HFModelSummary( repo_id="owner/oversized", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", file_size=100 * 1024**3, ) candidates, skipped, failed = process_model_for_candidates( model=model, hf_discovery=SafeDiscovery(), # type: ignore[arg-type] modelhub_client=OtherGpuProcessedClient(), # type: ignore[arg-type] template_selector=TemplateSelector(), target_gpus=["MetaX_c-500"], allowed_task_types=["text-generation"], preflight_advisor=CandidatePreflightAdvisor(), official_registry=registry, # type: ignore[arg-type] ) self.assertEqual([], candidates) self.assertEqual("preflight_predicted_oom", skipped[0]["reason"]) self.assertEqual([], failed) self.assertEqual(0, registry.calls) def test_lookup_failure_is_fail_closed(self) -> None: pool = ModelHubClientPool([FailingLookupClient()]) # type: ignore[arg-type] with self.assertRaises(ModelHubAPIError): pool.model_submission_precheck("owner/model") def test_submit_precheck_stops_an_exact_processed_gpu(self) -> None: client = ExactGpuProcessedClient() result = submit_candidate(candidate(), client) # type: ignore[arg-type] self.assertEqual("duplicate", result["outcome"]) self.assertEqual(0, client.add_calls) def test_uniqueness_rejection_is_non_retryable_for_that_combination(self) -> None: error = ModelHubAPIError("模型唯一性检查没有通过，无法进行同步") self.assertTrue(is_model_uniqueness_error(error)) result = submit_candidate(candidate(), UniquenessRejectingClient()) # type: ignore[arg-type] self.assertEqual("uniqueness_rejected", result["outcome"]) with tempfile.TemporaryDirectory() as temporary_dir: path = Path(temporary_dir) / "exclusions.jsonl" store = SubmissionExclusionStore(path) store.block("owner/model", "Vastai_va16", reason=str(error)) reloaded = SubmissionExclusionStore(path) self.assertTrue(reloaded.is_blocked("owner/model", "Vastai_va16")) self.assertFalse(reloaded.is_blocked("owner/model", "Biren_166m")) def test_run_persists_uniqueness_rejection_and_does_not_retry_it(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) args = build_parser().parse_args( [ "--gpus", "Vastai_va16", "--task-types", "text-generation", "--limit", "1", "--max-scan-models", "1", "--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.submission_exclusions_path = str(root / "exclusions.jsonl") args.history_archive_path = str(root / "history.jsonl") client = UniquenessRunClient() discovery = SafeDiscovery() first = run_submission( args, now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc), hf_discovery=discovery, # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] ) second = run_submission( args, now=datetime(2026, 1, 1, 12, 1, tzinfo=timezone.utc), hf_discovery=discovery, # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] ) self.assertEqual(1, first["modelGpuUniquenessRejectedCount"]) self.assertEqual(1, client.add_calls) self.assertEqual(0, second["candidateCount"]) self.assertEqual(1, second["skipReasonCounts"]["model_gpu_uniqueness_blocklist"]) def test_concurrent_batch_uses_at_most_one_gpu_per_model(self) -> None: selected = one_candidate_per_model( [ candidate("Vastai_va16"), candidate("Biren_166m"), {**candidate("Biren_166m"), "repoId": "owner/other"}, ] ) self.assertEqual(["owner/model", "owner/other"], [item["repoId"] for item in selected]) if __name__ == "__main__": unittest.main()