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 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 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_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_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()