Files
submmit/tests/test_concurrency.py
2026-08-11 00:59:40 +08:00

571 lines
22 KiB
Python

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