from __future__ import annotations import os import sys import tempfile import threading import unittest from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path from unittest.mock import patch 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 common import read_json, read_jsonl # noqa: E402 from main import build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402 from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error # noqa: E402 from models import HFModelSummary, ModelInspection # noqa: E402 from outcome_tracker import OutcomeTracker # noqa: E402 from submission_claims import SubmissionClaimStore, candidate_key # noqa: E402 from template_selector import TemplateSelector # noqa: E402 class FakeClient: def __init__(self, active_count: int = 0, *, reject_capacity: bool = False) -> None: self.active_count = active_count self.reject_capacity = reject_capacity self.submitted: list[dict] = [] self.count_calls = 0 self.search_calls = 0 self._lock = threading.Lock() def count_active_tasks(self, **_kwargs) -> int: # noqa: ANN003 with self._lock: self.count_calls += 1 return self.active_count def add_task(self, payload: dict) -> dict: if self.reject_capacity: raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限") with self._lock: self.submitted.append(payload) return {"code": 0, "data": {"id": len(self.submitted)}} def search_by_model_id(self, _model_id: str) -> dict: with self._lock: self.search_calls += 1 return {"code": 0, "data": {"verifyResult": {}}} def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003 return [] def list_tasks_page(self, **_kwargs) -> dict: # noqa: ANN003 return {"code": 0, "data": {"records": [], "pages": 0}} def find_recent_task_id(self, *_args, **_kwargs): # noqa: ANN002, ANN003 return None class FakeDiscovery: def __init__(self, count: int) -> None: self.models = [ HFModelSummary( repo_id=f"owner/model-{index}", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) for index in range(count) ] def list_recent_models(self, **_kwargs): # noqa: ANN003 return self.models @staticmethod def inspect_model(model: HFModelSummary) -> ModelInspection: return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"]) class DuplicateThenSuccessClient: def __init__(self) -> None: self.calls = 0 self.active = 0 @staticmethod def begin_cycle() -> None: return None def available_submit_slots(self) -> int: return max(0, 2 - self.active) @staticmethod def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003 return {"code": 0, "data": {"records": [], "pages": 0}} @staticmethod def processed_gpus_for_model(_model_id: str) -> set[str]: return set() def add_task(self, _payload: dict) -> dict: self.calls += 1 if self.calls <= 2: raise ModelHubAPIError("模型正在验证中，请勿重复提交") self.active += 1 return {"code": 0, "data": {"id": f"task-{self.calls}"}} class DynamicCapacityClient: def __init__(self, *, active: int, limit: int, token: str = "dynamic-test") -> None: self.active = active self.limit = limit self.token = token def count_active_tasks(self, *, max_count: int, **_kwargs) -> int: # noqa: ANN003 return min(self.active, max_count) def add_task(self, _payload: dict) -> dict: if self.active >= self.limit: raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限") self.active += 1 return {"code": 0, "data": {"id": self.active}} class AutoStrategyClient: def __init__(self, available: int) -> None: self.available = available self.calls = 0 @staticmethod def begin_cycle() -> None: return None def available_submit_slots(self) -> int: return self.available @staticmethod def list_tasks(**_kwargs) -> list[dict]: # noqa: ANN003 return [] @staticmethod def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003 return {"code": 0, "data": {"records": [], "pages": 0}} @staticmethod def processed_gpus_for_model(_model_id: str) -> set[str]: return set() def add_task(self, _payload: dict) -> dict: self.calls += 1 if self.calls == 1: raise ModelHubAPIError("模型正在验证中，请勿重复提交") self.available -= 1 return {"code": 0, "data": {"id": f"strategy-{self.calls}"}} class ExhaustedRecentDiscovery: def __init__(self) -> None: self.calls: list[object] = [] self.recent = HFModelSummary( repo_id="owner/recent-exhausted", downloads=100, last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) self.older = [ HFModelSummary( repo_id=f"owner/older-{index}", downloads=100, last_modified=datetime(2025, 1, 1, tzinfo=timezone.utc), pipeline_tag="text-generation", ) for index in range(3) ] def list_recent_models(self, *, updated_after=None, **_kwargs): # noqa: ANN003 self.calls.append(updated_after) return [self.recent] if updated_after is not None else [self.recent, *self.older] @staticmethod def inspect_model(model: HFModelSummary) -> ModelInspection: return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"]) class ExhaustedRecentClient(AutoStrategyClient): def __init__(self, available: int, processed_gpus: set[str]) -> None: super().__init__(available) self.processed_gpus = processed_gpus def processed_gpus_for_model(self, model_id: str) -> set[str]: if model_id == "owner/recent-exhausted": return set(self.processed_gpus) return set() def add_task(self, _payload: dict) -> dict: self.calls += 1 self.available -= 1 return {"code": 0, "data": {"id": f"fallback-{self.calls}"}} def make_candidate(index: int) -> dict: return { "repoId": f"owner/model-{index}", "modelAddress": f"https://modelscope.cn/models/owner/model-{index}", "targetGpu": "NVIDIA-A100", } class ClientPoolConcurrencyTests(unittest.TestCase): def test_online_submission_does_not_construct_llm_even_when_key_is_present(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) args = build_parser().parse_args( [ "--gpus", "Iluvatar_bi-150", "--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.history_archive_path = str(root / "history.jsonl") with ( patch.dict(os.environ, {"MODELHUB_QWEN_API_KEY": "must-not-be-used"}), patch( "llm_classifier.LLMAssistedClassifier.__init__", side_effect=AssertionError("online submission constructed an LLM client"), ), ): summary = run_submission( args, now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc), hf_discovery=FakeDiscovery(1), # type: ignore[arg-type] modelhub_client=AutoStrategyClient(available=1), # type: ignore[arg-type] ) self.assertFalse(summary["candidatePreflight"]["llm"]["enabled"]) def test_exhausted_recent_window_expands_to_older_models_in_same_run(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) args = build_parser().parse_args( [ "--task-types", "text-generation", "--since-hours", "48", "--limit", "4", "--max-scan-models", "4", "--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.history_archive_path = str(root / "history.jsonl") args.gpu_strategy_state_path = str(root / "strategy.json") selector = TemplateSelector() supported_gpus = set(selector.supported_target_gpus("text-generation", auto_only=True)) discovery = ExhaustedRecentDiscovery() client = ExhaustedRecentClient(available=2, processed_gpus=supported_gpus) summary = run_submission( args, now=datetime(2026, 1, 2, 12, tzinfo=timezone.utc), hf_discovery=discovery, # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] template_selector=selector, ) self.assertEqual(2, summary["submittedCount"]) self.assertEqual(len(supported_gpus), summary["skipReasonCounts"]["already_processed_for_gpu"]) self.assertEqual("all_history", summary["scanStages"][-1]["name"]) self.assertIsNone(discovery.calls[-1]) def test_adaptive_strategy_counts_only_platform_accepted_tasks(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) args = build_parser().parse_args( [ "--task-types", "text-generation", "--limit", "4", "--max-scan-models", "4", "--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.history_archive_path = str(root / "history.jsonl") args.gpu_strategy_state_path = str(root / "strategy.json") client = AutoStrategyClient(available=3) summary = run_submission( args, now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc), hf_discovery=FakeDiscovery(4), # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] ) state = read_json(root / "strategy.json") self.assertEqual(3, summary["submittedCount"]) self.assertEqual(1, summary["duplicateCount"]) self.assertEqual(3, state["acceptedSinceRefresh"]) self.assertEqual(3, sum(state["acceptedByCategory"].values())) def test_duplicate_batch_is_replaced_until_available_slots_are_filled(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) args = build_parser().parse_args( [ "--gpus", "Iluvatar_bi-150", "--task-types", "text-generation", "--limit", "4", "--max-scan-models", "4", "--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.history_archive_path = str(root / "history.jsonl") client = DuplicateThenSuccessClient() summary = run_submission( args, now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc), hf_discovery=FakeDiscovery(4), # type: ignore[arg-type] modelhub_client=client, # type: ignore[arg-type] ) self.assertEqual(2, summary["targetSubmitCount"]) self.assertEqual(4, summary["plannedSubmitCount"]) self.assertEqual(2, summary["duplicateCount"]) self.assertEqual(2, summary["submittedCount"]) self.assertEqual(0, client.available_submit_slots()) def test_platform_duplicate_is_classified_as_skipped_candidate(self) -> None: error = ModelHubAPIError("模型正在验证中，请勿重复提交") self.assertTrue(is_duplicate_submission_error(error)) class DuplicateClient: @staticmethod def add_task(_payload): # noqa: ANN001 raise error candidate = { **make_candidate(1), "taskType": "text-generation", "framework": "vllm", "configParams": "framework: vllm", } result = submit_candidate(candidate, DuplicateClient()) # type: ignore[arg-type] self.assertEqual("duplicate", result["outcome"]) def test_every_account_reaches_capacity_under_load(self) -> None: clients = [FakeClient() for _ in range(12)] pool = ModelHubClientPool( clients, # type: ignore[arg-type] active_task_cap=10, active_counts_ttl=60, instance_id="full-load-test", ) with ThreadPoolExecutor(max_workers=24) as executor: list(executor.map(lambda index: pool.add_task({"index": index}), range(120))) self.assertEqual([10] * 12, [len(client.submitted) for client in clients]) self.assertEqual([10] * 12, pool.active_task_counts()) self.assertEqual(0, pool.available_submit_slots()) def test_concurrent_submissions_reserve_and_balance_accounts(self) -> None: clients = [FakeClient(), FakeClient()] pool = ModelHubClientPool( clients, # type: ignore[arg-type] active_task_cap=4, active_counts_ttl=60, instance_id="balance-test", ) with ThreadPoolExecutor(max_workers=8) as executor: results = list(executor.map(lambda index: pool.add_task({"index": index}), range(8))) self.assertEqual(8, len(results)) self.assertEqual([4, 4], [len(client.submitted) for client in clients]) self.assertEqual([4, 4], pool.active_task_counts()) self.assertEqual(0, pool.available_submit_slots()) # A lagging remote count must not erase successful local reservations. pool._refresh_active_counts(force=True) self.assertEqual([4, 4], pool.active_task_counts()) def test_capacity_rejection_falls_through_to_another_account(self) -> None: full_elsewhere = FakeClient(active_count=0, reject_capacity=True) available = FakeClient(active_count=1) pool = ModelHubClientPool( [full_elsewhere, available], # type: ignore[arg-type] active_task_cap=2, instance_id="fallback-test", ) response = pool.add_task({"model": "x"}) self.assertEqual(1, response["data"]["id"]) self.assertEqual(0, len(full_elsewhere.submitted)) self.assertEqual(1, len(available.submitted)) def test_verification_cache_is_reused_between_cycles(self) -> None: client = FakeClient() pool = ModelHubClientPool([client], active_task_cap=2) # type: ignore[arg-type] pool.search_by_model_id("owner/model") pool.search_by_model_id("owner/model") self.assertEqual(1, client.search_calls) pool.begin_cycle() pool.search_by_model_id("owner/model") self.assertEqual(1, client.search_calls) def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None: client = DynamicCapacityClient(active=2, limit=3) pool = ModelHubClientPool( [client], # type: ignore[arg-type] active_task_cap=2, active_counts_ttl=60, capacity_probe_interval_cycles=3, instance_id="capacity-growth-test", ) pool.configure_capacity_probe(2) self.assertEqual(0, pool.available_submit_slots()) pool.configure_capacity_probe(3) self.assertEqual(1, pool.available_submit_slots()) pool.add_task({"model": "probe"}) self.assertEqual([3], pool.account_capacity_limits()) self.assertEqual([3], pool.active_task_counts()) self.assertEqual(0, pool.available_submit_slots()) def test_rejected_capacity_probe_keeps_known_limit_and_enters_cooldown(self) -> None: client = DynamicCapacityClient(active=2, limit=2) pool = ModelHubClientPool( [client], # type: ignore[arg-type] active_task_cap=2, active_counts_ttl=60, capacity_probe_interval_cycles=3, capacity_probe_cooldown_cycles=3, instance_id="capacity-rejection-test", ) pool.configure_capacity_probe(3) self.assertEqual(1, pool.available_submit_slots()) with self.assertRaises(ModelHubAPIError): pool.add_task({"model": "probe"}) self.assertEqual([2], pool.account_capacity_limits()) self.assertEqual(0, pool.available_submit_slots()) class ProcessCoordinationTests(unittest.TestCase): def test_local_failure_block_expires_after_24_hours(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: tracker = OutcomeTracker(Path(temporary_dir) / "outcomes.jsonl") now = datetime(2026, 1, 2, 12, tzinfo=timezone.utc) tracker._records = [ { "modelId": "owner/old-failure", "targetGpu": "gpu-a", "outcome": "failed", "submitTime": (now - timedelta(hours=26)).isoformat(), "lastSyncTime": (now - timedelta(hours=25)).isoformat(), }, { "modelId": "owner/recent-failure", "targetGpu": "gpu-a", "outcome": "failed", "submitTime": (now - timedelta(hours=2)).isoformat(), "lastSyncTime": (now - timedelta(hours=1)).isoformat(), }, ] tracker._rebuild_indexes() self.assertFalse(tracker.is_model_gpu_failed("owner/old-failure", "gpu-a", now=now)) self.assertTrue(tracker.is_model_gpu_failed("owner/recent-failure", "gpu-a", now=now)) def test_duplicate_claims_are_retained_so_next_batch_moves_forward(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: path = Path(temporary_dir) / "claims.jsonl" candidates = [make_candidate(index) for index in range(4)] store = SubmissionClaimStore(path, owner_id="worker-a") duplicates = store.claim(candidates, limit=2) store.mark_submitted(duplicates) replacements = store.claim(candidates, limit=2) self.assertEqual( {candidate_key(candidate) for candidate in candidates[2:]}, {candidate_key(candidate) for candidate in replacements}, ) def test_concurrent_claim_stores_select_disjoint_candidates(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: path = Path(temporary_dir) / "claims.jsonl" candidates = [make_candidate(index) for index in range(6)] first = SubmissionClaimStore(path, owner_id="worker-a") second = SubmissionClaimStore(path, owner_id="worker-b") with ThreadPoolExecutor(max_workers=2) as executor: first_future = executor.submit(first.claim, candidates, limit=3) second_future = executor.submit(second.claim, candidates, limit=3) first_claims = first_future.result() second_claims = second_future.result() first_keys = {candidate_key(candidate) for candidate in first_claims} second_keys = {candidate_key(candidate) for candidate in second_claims} self.assertEqual(3, len(first_claims)) self.assertEqual(3, len(second_claims)) self.assertTrue(first_keys.isdisjoint(second_keys)) def test_outcome_saves_merge_instead_of_overwriting_other_worker(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: path = Path(temporary_dir) / "outcomes.jsonl" first = OutcomeTracker(path) second = OutcomeTracker(path) first.record_submission("model-a", "gpu", "vllm", "text-generation", "task-a", "2026-01-01T00:00:00+00:00") second.record_submission("model-b", "gpu", "vllm", "text-generation", "task-b", "2026-01-01T00:00:01+00:00") first.save() second.save() self.assertEqual({"task-a", "task-b"}, {row["taskId"] for row in read_jsonl(path)}) def test_run_directories_are_allocated_atomically(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: base = Path(temporary_dir) / "runs" now = datetime(2026, 1, 1, tzinfo=timezone.utc) with ThreadPoolExecutor(max_workers=8) as executor: paths = list(executor.map(lambda _index: make_run_dir(base, now), range(8))) self.assertEqual(8, len(set(paths))) self.assertTrue(all(path.is_dir() for path in paths)) if __name__ == "__main__": unittest.main()